<commit_before>#![feature(asm)]

extern crate syscall;

use syscall::iopl;

use pci::{Pci, PciBar, PciClass};

mod pci;

fn enumerate_pci() {
    println!("PCI BS/DV/FN VEND:DEVI CL.SC.IN.RV");

    let pci = Pci::new();
    for bus in pci.buses() {
        for dev in bus.devs() {
            for func in dev.funcs() {
                if let Some(header) = func.header() {
                    println!("PCI {:>02X}/{:>02X}/{:>02X} {:>04X}:{:>04X} {:>02X}.{:>02X}.{:>02X}.{:>02X} {:?}",
                            bus.num, dev.num, func.num,
                            header.vendor_id, header.device_id,
                            header.class, header.subclass, header.interface, header.revision,
                            PciClass::from(header.class));

                    for i in 0..header.bars.len() {
                        match PciBar::from(header.bars[i]) {
                            PciBar::None => (),
                            PciBar::Memory(address) => println!("    BAR {} {:>08X}", i, address),
                            PciBar::Port(address) => println!("    BAR {} {:>04X}", i, address)
                        }
                    }

                    match PciClass::from(header.class) {
                        PciClass::Storage => match header.subclass {
                            0x01 => {
                                println!("    + IDE");
                            },
                            0x06 => {
                                println!("    + SATA");
                            },
                            _ => ()
                        },
                        PciClass::SerialBus => match header.subclass {
                            0x03 => match header.interface {
                                0x00 => {
                                    println!("    + UHCI");
                                },
                                0x10 => {
                                    println!("    + OHCI");
                                },
                                0x20 => {
                                    println!("    + EHCI");
                                },
                                0x30 => {
                                    println!("    + XHCI");
                                },
                                _ => ()
                            },
                            _ => ()
                        },
                        _ => ()
                    }
                }
            }
        }
    }
}

fn main() {
    unsafe { iopl(3).unwrap() };

    enumerate_pci();
}
<commit_msg>More compact output<commit_after>#![feature(asm)]

extern crate syscall;

use syscall::iopl;

use pci::{Pci, PciBar, PciClass};

mod pci;

fn enumerate_pci() {
    println!("PCI BS/DV/FN VEND:DEVI CL.SC.IN.RV");

    let pci = Pci::new();
    for bus in pci.buses() {
        for dev in bus.devs() {
            for func in dev.funcs() {
                if let Some(header) = func.header() {
                    print!("PCI {:>02X}/{:>02X}/{:>02X} {:>04X}:{:>04X} {:>02X}.{:>02X}.{:>02X}.{:>02X}",
                            bus.num, dev.num, func.num,
                            header.vendor_id, header.device_id,
                            header.class, header.subclass, header.interface, header.revision);

                    let pci_class = PciClass::from(header.class);
                    print!(" {:?}", pci_class);
                    match pci_class {
                        PciClass::Storage => match header.subclass {
                            0x01 => {
                                print!(" IDE");
                            },
                            0x06 => {
                                print!(" SATA");
                            },
                            _ => ()
                        },
                        PciClass::SerialBus => match header.subclass {
                            0x03 => match header.interface {
                                0x00 => {
                                    print!(" UHCI");
                                },
                                0x10 => {
                                    print!(" OHCI");
                                },
                                0x20 => {
                                    print!(" EHCI");
                                },
                                0x30 => {
                                    print!(" XHCI");
                                },
                                _ => ()
                            },
                            _ => ()
                        },
                        _ => ()
                    }

                    for i in 0..header.bars.len() {
                        match PciBar::from(header.bars[i]) {
                            PciBar::None => (),
                            PciBar::Memory(address) => print!(" {}={:>08X}", i, address),
                            PciBar::Port(address) => print!(" {}={:>04X}", i, address)
                        }
                    }

                    print!("\n");
                }
            }
        }
    }
}

fn main() {
    unsafe { iopl(3).unwrap() };

    enumerate_pci();
}
<|endoftext|>


<commit_before>use std::cmp;

use ctxt::SemContext;

use gc::Address;
use gc::root::IndirectObj;
use gc::swiper::CARD_SIZE;
use gc::swiper::card::{CardEntry, CardTable};
use gc::swiper::crossing::{Card, CrossingEntry, CrossingMap};
use gc::swiper::old::OldGen;
use gc::swiper::young::YoungGen;

use mem;
use object::Obj;
use os::{self, ProtType};
use timer::{in_ms, Timer};

pub struct MinorCollector<'a, 'ast: 'a> {
    ctxt: &'a SemContext<'ast>,
    young: &'a YoungGen,
    old: &'a OldGen,
    rootset: &'a [IndirectObj],
    card_table: &'a CardTable,
    crossing_map: &'a CrossingMap,

    free: Address,
    promotion_failed: bool,
}

impl<'a, 'ast> MinorCollector<'a, 'ast> {
    pub fn new(
        ctxt: &'a SemContext<'ast>,
        young: &'a YoungGen,
        old: &'a OldGen,
        card_table: &'a CardTable,
        crossing_map: &'a CrossingMap,
        rootset: &'a [IndirectObj],
    ) -> MinorCollector<'a, 'ast> {
        MinorCollector {
            ctxt: ctxt,
            young: young,
            old: old,
            rootset: rootset,
            card_table: card_table,
            crossing_map: crossing_map,
            free: Address::null(),
            promotion_failed: false,
        }
    }

    pub fn collect(&mut self) {
        let mut timer = Timer::new(self.ctxt.args.flag_gc_events);
        let to_space = self.young.to_space();

        let mut scan = to_space.start;
        self.free = scan;

        if cfg!(debug_assertions) {
            // make memory readable & writable again, so that we
            // can copy objects to the to-space.
            // Since this has some overhead, do it only in debug builds.
            os::mprotect(
                to_space.start.to_ptr::<u8>(),
                to_space.size(),
                ProtType::Writable,
            );
        }

        // detect all references from roots into young generation
        for &root in self.rootset {
            let root_ptr = root.get();

            if self.young.contains(Address::from_ptr(root_ptr)) {
                root.set(self.copy(root_ptr));
            }
        }

        // detect references from old generation (dirty cards) into young generation
        self.copy_dirty_cards();

        // visit all fields in copied objects
        while scan < self.free {
            let object = unsafe { &mut *scan.to_mut_ptr::<Obj>() };

            object.visit_reference_fields(|field| {
                let field_ptr = field.get();

                if self.young.contains(Address::from_ptr(field_ptr)) {
                    field.set(self.copy(field_ptr));
                }
            });

            scan = scan.offset(object.size());
        }

        self.young.swap_spaces(self.free);

        // Make from-space unaccessible both from read/write.
        // Since this has some overhead, do it only in debug builds.
        if cfg!(debug_assertions) {
            let to_space = self.young.to_space();
            os::mprotect(
                to_space.start.to_ptr::<u8>(),
                to_space.size(),
                ProtType::None,
            );
        }

        timer.stop_with(|dur| if self.ctxt.args.flag_gc_events {
            println!("GC minor: collect garbage ({} ms)", in_ms(dur));
        });
    }

    // copy all references from old- into young-generation.
    fn copy_dirty_cards(&mut self) {
        self.card_table.visit_dirty(|card| {
            let crossing_entry = self.crossing_map.get(card);
            let card_start = self.old.address_from_card(card);
            let card_end = card_start.offset(CARD_SIZE);

            match crossing_entry {
                CrossingEntry::NoRefs => panic!("card dirty without any refs"),
                CrossingEntry::LeadingRefs(refs) => {
                    let mut ptr = card_start;

                    for _ in 0..refs {
                        let ind_ptr = IndirectObj::from_address(ptr);
                        let dir_ptr = ind_ptr.get();

                        if self.young.contains(Address::from_ptr(dir_ptr)) {
                            ind_ptr.set(self.copy(dir_ptr));
                        }

                        ptr = ptr.offset(mem::ptr_width() as usize);
                    }

                    // copy all objects from this card
                    self.copy_card(card, ptr, card_end);
                }

                CrossingEntry::FirstObjectOffset(offset) => {
                    let ptr = card_start.offset(offset as usize * mem::ptr_width_usize());

                    // copy all objects from this card
                    self.copy_card(card, ptr, card_end);
                }
            }
        });
    }

    fn copy_card(&mut self, card: Card, mut ptr: Address, card_end: Address) {
        let old_end: Address = self.old.free();
        let mut end = cmp::min(card_end, old_end);
        let mut ref_to_young_gen = false;

        loop {
            while ptr < end {
                let object = unsafe { &mut *ptr.to_mut_ptr::<Obj>() };

                object.visit_reference_fields(|field| {
                    let field_ptr = field.get();

                    if self.young.contains(Address::from_ptr(field_ptr)) {
                        let copied_obj = self.copy(field_ptr);
                        field.set(copied_obj);

                        // determine if copied object is still in young generation
                        if self.young.contains(Address::from_ptr(copied_obj)) {
                            ref_to_young_gen = true;
                        }
                    }
                });

                ptr = ptr.offset(object.size());
            }

            // if we are in the last card of the old generation, promoted objects
            // will increase `end` towards `card_end`. Those newly promoted objects
            // need also be handled.
            if end == card_end {
                break;
            }

            let old_end = self.old.free();
            let next_end = cmp::min(card_end, old_end);

            if end != next_end {
                end = next_end;
            } else {
                break;
            }
        }

        // if there are no references to the young generation in this card,
        // set the card to clean.
        if !ref_to_young_gen {
            self.card_table.set(card, CardEntry::Clean);
        }
    }

    fn copy(&mut self, obj: *mut Obj) -> *mut Obj {
        let obj_addr = Address::from_ptr(obj);
        let obj = unsafe { &mut *obj };

        if let Some(fwd) = obj.header().forwarded() {
            return fwd.to_mut_ptr();
        }

        let obj_size = obj.size();

        // if object is old enough we copy it into the old generation
        if self.young.should_be_promoted(obj_addr) {
            let copy_addr = Address::from_ptr(self.old.alloc(obj_size));

            // if there isn't enough space in old gen keep it in the
            // young generation for now
            if copy_addr.is_null() {
                self.promotion_failed = true;

            } else {
                self.promote_object(obj, copy_addr, obj_size);
                return copy_addr.to_mut_ptr();
            }
        }

        copy_addr = self.free;
        self.free = copy_addr.offset(obj_size);

        obj.copy_to(copy_addr, obj_size);
        obj.header_mut().forward_to(copy_addr);

        copy_addr.to_mut_ptr()
    }

    fn promote_object(&mut self, obj: &mut Obj, copy_addr: Address, obj_size: usize) {
        obj.copy_to(copy_addr, obj_size);

        // Promoted object can have references to the young generation.
        // Set the card table entry to dirty if this is the case.
        self.handle_promoted_object(copy_addr);

        obj.header_mut().forward_to(copy_addr);
    }

    fn handle_promoted_object(&mut self, addr: Address) {
        let card = self.old.card_from_address(addr);

        // card is already dirty, nothing left to do.
        if self.card_table.get(card).is_dirty() {
            return;
        }

        let object: &mut Obj = unsafe { &mut *addr.to_mut_ptr() };
        let mut old_to_young_ref = false;

        object.visit_reference_fields(|field| {
            let field_ptr = field.get();

            if self.young.contains(Address::from_ptr(field_ptr)) {
                old_to_young_ref = true;
            }
        });

        if old_to_young_ref {
            self.card_table.set(card, CardEntry::Dirty);
        }
    }

    pub fn promotion_failed(&self) -> bool {
        self.promotion_failed
    }
}
<commit_msg>fix compile error<commit_after>use std::cmp;

use ctxt::SemContext;

use gc::Address;
use gc::root::IndirectObj;
use gc::swiper::CARD_SIZE;
use gc::swiper::card::{CardEntry, CardTable};
use gc::swiper::crossing::{Card, CrossingEntry, CrossingMap};
use gc::swiper::old::OldGen;
use gc::swiper::young::YoungGen;

use mem;
use object::Obj;
use os::{self, ProtType};
use timer::{in_ms, Timer};

pub struct MinorCollector<'a, 'ast: 'a> {
    ctxt: &'a SemContext<'ast>,
    young: &'a YoungGen,
    old: &'a OldGen,
    rootset: &'a [IndirectObj],
    card_table: &'a CardTable,
    crossing_map: &'a CrossingMap,

    free: Address,
    promotion_failed: bool,
}

impl<'a, 'ast> MinorCollector<'a, 'ast> {
    pub fn new(
        ctxt: &'a SemContext<'ast>,
        young: &'a YoungGen,
        old: &'a OldGen,
        card_table: &'a CardTable,
        crossing_map: &'a CrossingMap,
        rootset: &'a [IndirectObj],
    ) -> MinorCollector<'a, 'ast> {
        MinorCollector {
            ctxt: ctxt,
            young: young,
            old: old,
            rootset: rootset,
            card_table: card_table,
            crossing_map: crossing_map,
            free: Address::null(),
            promotion_failed: false,
        }
    }

    pub fn collect(&mut self) {
        let mut timer = Timer::new(self.ctxt.args.flag_gc_events);
        let to_space = self.young.to_space();

        let mut scan = to_space.start;
        self.free = scan;

        if cfg!(debug_assertions) {
            // make memory readable & writable again, so that we
            // can copy objects to the to-space.
            // Since this has some overhead, do it only in debug builds.
            os::mprotect(
                to_space.start.to_ptr::<u8>(),
                to_space.size(),
                ProtType::Writable,
            );
        }

        // detect all references from roots into young generation
        for &root in self.rootset {
            let root_ptr = root.get();

            if self.young.contains(Address::from_ptr(root_ptr)) {
                root.set(self.copy(root_ptr));
            }
        }

        // detect references from old generation (dirty cards) into young generation
        self.copy_dirty_cards();

        // visit all fields in copied objects
        while scan < self.free {
            let object = unsafe { &mut *scan.to_mut_ptr::<Obj>() };

            object.visit_reference_fields(|field| {
                let field_ptr = field.get();

                if self.young.contains(Address::from_ptr(field_ptr)) {
                    field.set(self.copy(field_ptr));
                }
            });

            scan = scan.offset(object.size());
        }

        self.young.swap_spaces(self.free);

        // Make from-space unaccessible both from read/write.
        // Since this has some overhead, do it only in debug builds.
        if cfg!(debug_assertions) {
            let to_space = self.young.to_space();
            os::mprotect(
                to_space.start.to_ptr::<u8>(),
                to_space.size(),
                ProtType::None,
            );
        }

        timer.stop_with(|dur| if self.ctxt.args.flag_gc_events {
            println!("GC minor: collect garbage ({} ms)", in_ms(dur));
        });
    }

    // copy all references from old- into young-generation.
    fn copy_dirty_cards(&mut self) {
        self.card_table.visit_dirty(|card| {
            let crossing_entry = self.crossing_map.get(card);
            let card_start = self.old.address_from_card(card);
            let card_end = card_start.offset(CARD_SIZE);

            match crossing_entry {
                CrossingEntry::NoRefs => panic!("card dirty without any refs"),
                CrossingEntry::LeadingRefs(refs) => {
                    let mut ptr = card_start;

                    for _ in 0..refs {
                        let ind_ptr = IndirectObj::from_address(ptr);
                        let dir_ptr = ind_ptr.get();

                        if self.young.contains(Address::from_ptr(dir_ptr)) {
                            ind_ptr.set(self.copy(dir_ptr));
                        }

                        ptr = ptr.offset(mem::ptr_width() as usize);
                    }

                    // copy all objects from this card
                    self.copy_card(card, ptr, card_end);
                }

                CrossingEntry::FirstObjectOffset(offset) => {
                    let ptr = card_start.offset(offset as usize * mem::ptr_width_usize());

                    // copy all objects from this card
                    self.copy_card(card, ptr, card_end);
                }
            }
        });
    }

    fn copy_card(&mut self, card: Card, mut ptr: Address, card_end: Address) {
        let old_end: Address = self.old.free();
        let mut end = cmp::min(card_end, old_end);
        let mut ref_to_young_gen = false;

        loop {
            while ptr < end {
                let object = unsafe { &mut *ptr.to_mut_ptr::<Obj>() };

                object.visit_reference_fields(|field| {
                    let field_ptr = field.get();

                    if self.young.contains(Address::from_ptr(field_ptr)) {
                        let copied_obj = self.copy(field_ptr);
                        field.set(copied_obj);

                        // determine if copied object is still in young generation
                        if self.young.contains(Address::from_ptr(copied_obj)) {
                            ref_to_young_gen = true;
                        }
                    }
                });

                ptr = ptr.offset(object.size());
            }

            // if we are in the last card of the old generation, promoted objects
            // will increase `end` towards `card_end`. Those newly promoted objects
            // need also be handled.
            if end == card_end {
                break;
            }

            let old_end = self.old.free();
            let next_end = cmp::min(card_end, old_end);

            if end != next_end {
                end = next_end;
            } else {
                break;
            }
        }

        // if there are no references to the young generation in this card,
        // set the card to clean.
        if !ref_to_young_gen {
            self.card_table.set(card, CardEntry::Clean);
        }
    }

    fn copy(&mut self, obj: *mut Obj) -> *mut Obj {
        let obj_addr = Address::from_ptr(obj);
        let obj = unsafe { &mut *obj };

        if let Some(fwd) = obj.header().forwarded() {
            return fwd.to_mut_ptr();
        }

        let obj_size = obj.size();

        // if object is old enough we copy it into the old generation
        if self.young.should_be_promoted(obj_addr) {
            let copy_addr = Address::from_ptr(self.old.alloc(obj_size));

            // if there isn't enough space in old gen keep it in the
            // young generation for now
            if copy_addr.is_null() {
                self.promotion_failed = true;

            } else {
                self.promote_object(obj, copy_addr, obj_size);
                return copy_addr.to_mut_ptr();
            }
        }

        let copy_addr = self.free;
        self.free = copy_addr.offset(obj_size);

        obj.copy_to(copy_addr, obj_size);
        obj.header_mut().forward_to(copy_addr);

        copy_addr.to_mut_ptr()
    }

    fn promote_object(&mut self, obj: &mut Obj, copy_addr: Address, obj_size: usize) {
        obj.copy_to(copy_addr, obj_size);

        // Promoted object can have references to the young generation.
        // Set the card table entry to dirty if this is the case.
        self.handle_promoted_object(copy_addr);

        obj.header_mut().forward_to(copy_addr);
    }

    fn handle_promoted_object(&mut self, addr: Address) {
        let card = self.old.card_from_address(addr);

        // card is already dirty, nothing left to do.
        if self.card_table.get(card).is_dirty() {
            return;
        }

        let object: &mut Obj = unsafe { &mut *addr.to_mut_ptr() };
        let mut old_to_young_ref = false;

        object.visit_reference_fields(|field| {
            let field_ptr = field.get();

            if self.young.contains(Address::from_ptr(field_ptr)) {
                old_to_young_ref = true;
            }
        });

        if old_to_young_ref {
            self.card_table.set(card, CardEntry::Dirty);
        }
    }

    pub fn promotion_failed(&self) -> bool {
        self.promotion_failed
    }
}
<|endoftext|>


<commit_before>use bytes::IntoBuf;
use futures::{Async, Future, Poll, Stream};
use futures::future::{self, Either};
use futures::sync::oneshot;
use h2::client::{Builder, Handshake, SendRequest};
use tokio_io::{AsyncRead, AsyncWrite};

use body::Payload;
use ::common::{Exec, Never};
use super::{PipeToSendStream, SendBuf};
use ::{Body, Request, Response};

type ClientRx<B> = ::client::dispatch::Receiver<Request<B>, Response<Body>>;

pub struct Client<T, B>
where
    B: Payload,
{
    executor: Exec,
    rx: ClientRx<B>,
    state: State<T, SendBuf<B::Data>>,
}

enum State<T, B> where B: IntoBuf {
    Handshaking(Handshake<T, B>),
    Ready(SendRequest<B>, oneshot::Sender<Never>),
}

impl<T, B> Client<T, B>
where
    T: AsyncRead + AsyncWrite + Send + 'static,
    B: Payload,
{
    pub(crate) fn new(io: T, rx: ClientRx<B>, exec: Exec) -> Client<T, B> {
        let handshake = Builder::new()
            // we don't expose PUSH promises yet
            .enable_push(false)
            .handshake(io);

        Client {
            executor: exec,
            rx: rx,
            state: State::Handshaking(handshake),
        }
    }
}

impl<T, B> Future for Client<T, B>
where
    T: AsyncRead + AsyncWrite + Send + 'static,
    B: Payload + 'static,
{
    type Item = ();
    type Error = ::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let next = match self.state {
                State::Handshaking(ref mut h) => {
                    let (request_tx, conn) = try_ready!(h.poll().map_err(::Error::new_h2));
                    // A oneshot channel is used entirely to detect when the
                    // 'Client' has been dropped. This is to get around a bug
                    // in h2 where dropping all SendRequests won't notify a
                    // parked Connection.
                    let (tx, rx) = oneshot::channel();
                    let fut = conn
                        .inspect(|_| trace!("connection complete"))
                        .map_err(|e| debug!("connection error: {}", e))
                        .select2(rx)
                        .then(|res| match res {
                            Ok(Either::A(((), _))) |
                            Err(Either::A(((), _))) => {
                                // conn has finished either way
                                Either::A(future::ok(()))
                            },
                            Err(Either::B((_, conn))) => {
                                // oneshot has been dropped, hopefully polling
                                // the connection some more should start shutdown
                                // and then close
                                trace!("send_request dropped, starting conn shutdown");
                                Either::B(conn)
                            }
                            Ok(Either::B((never, _))) => match never {},
                        });
                    self.executor.execute(fut);
                    State::Ready(request_tx, tx)
                },
                State::Ready(ref mut tx, _) => {
                    try_ready!(tx.poll_ready().map_err(::Error::new_h2));
                    match self.rx.poll() {
                        Ok(Async::Ready(Some((req, mut cb)))) => {
                            // check that future hasn't been canceled already
                            if let Async::Ready(()) = cb.poll_cancel().expect("poll_cancel cannot error") {
                                trace!("request canceled");
                                continue;
                            }
                            let (head, body) = req.into_parts();
                            let mut req = ::http::Request::from_parts(head, ());
                            super::strip_connection_headers(req.headers_mut());
                            let eos = body.is_end_stream();
                            let (fut, body_tx) = match tx.send_request(req, eos) {
                                Ok(ok) => ok,
                                Err(err) => {
                                    debug!("client send request error: {}", err);
                                    let _ = cb.send(Err((::Error::new_h2(err), None)));
                                    continue;
                                }
                            };
                            if !eos {
                                let pipe = PipeToSendStream::new(body, body_tx);
                                self.executor.execute(pipe.map_err(|e| debug!("client request body error: {}", e)));
                            }

                            let fut = fut
                                .then(move |result| {
                                    match result {
                                        Ok(res) => {
                                            let res = res.map(::Body::h2);
                                            let _ = cb.send(Ok(res));
                                        },
                                        Err(err) => {
                                            debug!("client response error: {}", err);
                                            let _ = cb.send(Err((::Error::new_h2(err), None)));
                                        }
                                    }
                                    Ok(())
                                });
                            self.executor.execute(fut);
                            continue;
                        },

                        Ok(Async::NotReady) => return Ok(Async::NotReady),

                        Ok(Async::Ready(None)) |
                        Err(_) => {
                            trace!("client::dispatch::Sender dropped");
                            return Ok(Async::Ready(()));
                        }
                    }
                },
            };
            self.state = next;
        }
    }
}
<commit_msg>fix(http2): force notify h2 client connection when all body streams drop<commit_after>use bytes::IntoBuf;
use futures::{Async, Future, Poll, Stream};
use futures::future::{self, Either};
use futures::sync::mpsc;
use h2::client::{Builder, Handshake, SendRequest};
use tokio_io::{AsyncRead, AsyncWrite};

use body::Payload;
use ::common::{Exec, Never};
use super::{PipeToSendStream, SendBuf};
use ::{Body, Request, Response};

type ClientRx<B> = ::client::dispatch::Receiver<Request<B>, Response<Body>>;
/// An mpsc channel is used to help notify the `Connection` task when *all*
/// other handles to it have been dropped, so that it can shutdown.
type ConnDropRef = mpsc::Sender<Never>;

pub struct Client<T, B>
where
    B: Payload,
{
    executor: Exec,
    rx: ClientRx<B>,
    state: State<T, SendBuf<B::Data>>,
}

enum State<T, B> where B: IntoBuf {
    Handshaking(Handshake<T, B>),
    Ready(SendRequest<B>, ConnDropRef),
}

impl<T, B> Client<T, B>
where
    T: AsyncRead + AsyncWrite + Send + 'static,
    B: Payload,
{
    pub(crate) fn new(io: T, rx: ClientRx<B>, exec: Exec) -> Client<T, B> {
        let handshake = Builder::new()
            // we don't expose PUSH promises yet
            .enable_push(false)
            .handshake(io);

        Client {
            executor: exec,
            rx: rx,
            state: State::Handshaking(handshake),
        }
    }
}

impl<T, B> Future for Client<T, B>
where
    T: AsyncRead + AsyncWrite + Send + 'static,
    B: Payload + 'static,
{
    type Item = ();
    type Error = ::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let next = match self.state {
                State::Handshaking(ref mut h) => {
                    let (request_tx, conn) = try_ready!(h.poll().map_err(::Error::new_h2));
                    // An mpsc channel is used entirely to detect when the
                    // 'Client' has been dropped. This is to get around a bug
                    // in h2 where dropping all SendRequests won't notify a
                    // parked Connection.
                    let (tx, rx) = mpsc::channel(0);
                    let rx = rx.into_future()
                        .map(|(msg, _)| match msg {
                            Some(never) => match never {},
                            None => (),
                        })
                        .map_err(|_| -> Never { unreachable!("mpsc cannot error") });
                    let fut = conn
                        .inspect(|_| trace!("connection complete"))
                        .map_err(|e| debug!("connection error: {}", e))
                        .select2(rx)
                        .then(|res| match res {
                            Ok(Either::A(((), _))) |
                            Err(Either::A(((), _))) => {
                                // conn has finished either way
                                Either::A(future::ok(()))
                            },
                            Ok(Either::B(((), conn))) => {
                                // mpsc has been dropped, hopefully polling
                                // the connection some more should start shutdown
                                // and then close
                                trace!("send_request dropped, starting conn shutdown");
                                Either::B(conn)
                            }
                            Err(Either::B((never, _))) => match never {},
                        });
                    self.executor.execute(fut);
                    State::Ready(request_tx, tx)
                },
                State::Ready(ref mut tx, ref conn_dropper) => {
                    try_ready!(tx.poll_ready().map_err(::Error::new_h2));
                    match self.rx.poll() {
                        Ok(Async::Ready(Some((req, mut cb)))) => {
                            // check that future hasn't been canceled already
                            if let Async::Ready(()) = cb.poll_cancel().expect("poll_cancel cannot error") {
                                trace!("request canceled");
                                continue;
                            }
                            let (head, body) = req.into_parts();
                            let mut req = ::http::Request::from_parts(head, ());
                            super::strip_connection_headers(req.headers_mut());
                            let eos = body.is_end_stream();
                            let (fut, body_tx) = match tx.send_request(req, eos) {
                                Ok(ok) => ok,
                                Err(err) => {
                                    debug!("client send request error: {}", err);
                                    let _ = cb.send(Err((::Error::new_h2(err), None)));
                                    continue;
                                }
                            };
                            if !eos {
                                let conn_drop_ref = conn_dropper.clone();
                                let pipe = PipeToSendStream::new(body, body_tx)
                                    .map_err(|e| debug!("client request body error: {}", e))
                                    .then(move |x| {
                                        drop(conn_drop_ref);
                                        x
                                    });
                                self.executor.execute(pipe);
                            }

                            let fut = fut
                                .then(move |result| {
                                    match result {
                                        Ok(res) => {
                                            let res = res.map(::Body::h2);
                                            let _ = cb.send(Ok(res));
                                        },
                                        Err(err) => {
                                            debug!("client response error: {}", err);
                                            let _ = cb.send(Err((::Error::new_h2(err), None)));
                                        }
                                    }
                                    Ok(())
                                });
                            self.executor.execute(fut);
                            continue;
                        },

                        Ok(Async::NotReady) => return Ok(Async::NotReady),

                        Ok(Async::Ready(None)) |
                        Err(_) => {
                            trace!("client::dispatch::Sender dropped");
                            return Ok(Async::Ready(()));
                        }
                    }
                },
            };
            self.state = next;
        }
    }
}
<|endoftext|>


<commit_before>//! UDP relay proxy server

use std::{
    io::{self, Cursor},
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use bytes::BytesMut;
use futures::{future, future::AbortHandle, stream::FuturesUnordered, StreamExt};
use log::{debug, error, info, trace, warn};
use lru_time_cache::{Entry, LruCache};
use tokio::{
    self,
    net::udp::{RecvHalf, SendHalf},
    sync::{mpsc, Mutex},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        flow::{SharedMultiServerFlowStatistic, SharedServerFlowStatistic},
        socks5::Address,
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

use super::{
    crypto_io::{decrypt_payload, encrypt_payload},
    DEFAULT_TIMEOUT,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

// Represent a UDP association
struct UdpAssociation {
    // local -> remote Queue
    // Drops tx, will close local -> remote task
    tx: mpsc::Sender<Vec<u8>>,

    // local <- remote task life watcher
    watcher: AbortHandle,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.watcher.abort();
    }
}

impl UdpAssociation {
    /// Create an association with addr
    async fn associate(
        context: SharedContext,
        svr_idx: usize,
        src_addr: SocketAddr,
        mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
    ) -> io::Result<UdpAssociation> {
        // Create a socket for receiving packets
        let local_addr = match context.config().local_addr {
            None => {
                // Let system allocate an address for us
                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
            }
            Some(ref addr) => {
                // Uses configured local address
                addr.bind_addr(&context).await?
            }
        };
        let remote_udp = create_udp_socket(&local_addr).await?;

        let local_addr = remote_udp.local_addr().expect("could not determine port bound to");
        debug!("created UDP Association for {} from {}", src_addr, local_addr);

        // Create a channel for sending packets to remote
        // FIXME: Channel size 1024?
        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(1024);

        // Splits socket into sender and receiver
        let (mut receiver, mut sender) = remote_udp.split();

        let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);

        // local -> remote
        {
            let context = context.clone();
            tokio::spawn(async move {
                let svr_cfg = context.server_config(svr_idx);

                while let Some(pkt) = rx.recv().await {
                    // pkt is already a raw packet, so just send it
                    if let Err(err) =
                        UdpAssociation::relay_l2r(&context, src_addr, &mut sender, &pkt[..], timeout, svr_cfg).await
                    {
                        error!("failed to relay packet, {} -> ..., error: {}", src_addr, err);

                        // FIXME: Ignore? Or how to deal with it?
                    }
                }

                debug!("UDP ASSOCIATE {} -> .. finished", src_addr);
            });
        }

        let (r2l_task, close_flag) = future::abortable(async move {
            let svr_cfg = context.server_config(svr_idx);

            loop {
                // Read and send back to source
                match UdpAssociation::relay_r2l(&context, src_addr, &mut receiver, &mut response_tx, svr_cfg).await {
                    Ok(..) => {}
                    Err(err) => {
                        error!("failed to receive packet, {} <- .., error: {}", src_addr, err);

                        // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                        // break;
                    }
                }
            }
        });

        // local <- remote
        tokio::spawn(async move {
            let _ = r2l_task.await;

            debug!("UDP ASSOCIATE {} <- .. finished", src_addr);
        });

        Ok(UdpAssociation {
            tx,
            watcher: close_flag,
        })
    }

    /// Relay packets from local to remote
    async fn relay_l2r(
        context: &Context,
        src: SocketAddr,
        remote_udp: &mut SendHalf,
        pkt: &[u8],
        timeout: Duration,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // First of all, decrypt payload CLIENT -> SERVER
        let decrypted_pkt = match decrypt_payload(context, svr_cfg.method(), svr_cfg.key(), pkt) {
            Ok(Some(pkt)) => pkt,
            Ok(None) => {
                error!("failed to decrypt pkt in UDP relay, packet too short");
                let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
                return Err(err);
            }
            Err(err) => {
                error!("failed to decrypt pkt in UDP relay: {}", err);
                let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed");
                return Err(err);
            }
        };

        // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
        let mut cur = Cursor::new(decrypted_pkt);

        let addr = Address::read_from(&mut cur).await?;

        debug!("UDP ASSOCIATE {} <-> {} establishing", src, addr);

        if context.check_outbound_blocked(&addr) {
            warn!("outbound {} is blocked by ACL rules", addr);
            return Ok(());
        }

        // Take out internal buffer for optimizing one byte copy
        let header_len = cur.position() as usize;
        let decrypted_pkt = cur.into_inner();
        let body = &decrypted_pkt[header_len..];

        let send_len = match addr {
            Address::SocketAddress(ref remote_addr) => {
                debug!(
                    "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                    src,
                    addr,
                    remote_addr,
                    body.len()
                );
                try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await?
            }
            Address::DomainNameAddress(ref dname, port) => lookup_outbound_then!(context, dname, port, |remote_addr| {
                match try_timeout(remote_udp.send_to(body, &remote_addr), Some(timeout)).await {
                    Ok(l) => {
                        debug!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Ok(l)
                    }
                    Err(err) => {
                        error!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Err(err)
                    }
                }
            })
            .map(|(_, l)| l)?,
        };

        assert_eq!(body.len(), send_len);

        Ok(())
    }

    /// Relay packets from remote to local
    async fn relay_r2l(
        context: &Context,
        src_addr: SocketAddr,
        remote_udp: &mut RecvHalf,
        response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // Waiting for response from server SERVER -> CLIENT
        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let mut remote_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        let (remote_recv_len, remote_addr) = remote_udp.recv_from(&mut remote_buf).await?;

        debug!(
            "UDP ASSOCIATE {} <- {}, payload length {} bytes",
            src_addr, remote_addr, remote_recv_len
        );

        // FIXME: The Address should be the Address that client sent
        let addr = Address::SocketAddress(remote_addr);

        // CLIENT <- SERVER protocol: ADDRESS + PAYLOAD
        let mut send_buf = Vec::new();
        addr.write_to_buf(&mut send_buf);
        send_buf.extend_from_slice(&remote_buf[..remote_recv_len]);

        let mut encrypt_buf = BytesMut::new();
        encrypt_payload(context, svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?;

        // Send back to src_addr
        if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await {
            error!("failed to send packet into response channel, error: {}", err);

            // FIXME: What to do? Ignore?
        }

        Ok(())
    }

    // Send packet to remote
    //
    // Return `Err` if receiver have been closed
    async fn send(&mut self, pkt: Vec<u8>) {
        if let Err(..) = self.tx.send(pkt).await {
            // SHOULDn't HAPPEN
            unreachable!("UDP Association local -> remote Queue closed unexpectly");
        }
    }
}

async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, svr_idx: usize) -> io::Result<()> {
    let svr_cfg = context.server_config(svr_idx);
    let listen_addr = svr_cfg.addr().bind_addr(&context).await?;

    let listener = create_udp_socket(&listen_addr).await?;
    let local_addr = listener.local_addr().expect("determine port bound to");
    info!("shadowsocks UDP listening on {}", local_addr);

    let (mut r, mut w) = listener.split();

    // NOTE: Associations are only eliminated by expire time by default
    // So it may exhaust all available file descriptors
    let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);
    let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations {
        LruCache::with_expiry_duration_and_capacity(timeout, max_assoc)
    } else {
        LruCache::with_expiry_duration(timeout)
    };
    let assoc_map = Arc::new(Mutex::new(assoc_map));

    // FIXME: Channel size 1024?
    let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024);

    {
        // Tokio task for sending data back to clients

        let assoc_map = assoc_map.clone();
        let flow_stat = flow_stat.clone();

        tokio::spawn(async move {
            while let Some((src, pkt)) = rx.recv().await {
                let cache_key = src.to_string();
                {
                    let mut amap = assoc_map.lock().await;

                    // Check or update expire time
                    if amap.get(&cache_key).is_none() {
                        debug!(
                            "UDP association {} <-> ... is already expired, throwing away packet {} bytes",
                            src,
                            pkt.len()
                        );
                        continue;
                    }
                }

                if let Err(err) = w.send_to(&pkt, &src).await {
                    error!("UDP packet send failed, err: {:?}", err);
                    break;
                }

                flow_stat.udp().incr_tx(pkt.len() as u64);
            }

            // FIXME: How to stop the outer listener Future?
        });
    }

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src) = r.recv_from(&mut pkt_buf).await?;

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let pkt = &pkt_buf[..recv_len];

        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
        flow_stat.udp().incr_rx(pkt.len() as u64);

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check ACL
        if context.check_client_blocked(&src) {
            warn!("client {} is blocked by ACL rules", src);
            continue;
        }

        // Check or (re)create an association
        {
            // Locks the whole association map
            let mut assoc_map = assoc_map.lock().await;

            // Get or create an association
            let assoc = match assoc_map.entry(src.to_string()) {
                Entry::Occupied(oc) => oc.into_mut(),
                Entry::Vacant(vc) => vc.insert(
                    UdpAssociation::associate(context.clone(), svr_idx, src, tx.clone())
                        .await
                        .expect("create udp association"),
                ),
            };

            // FIXME: Lock is still kept for a mutable reference
            // Send to local -> remote task
            assoc.send(pkt.to_vec()).await;
        }
    }
}

/// Starts a UDP relay server
pub async fn run(context: SharedContext, flow_stat: SharedMultiServerFlowStatistic) -> io::Result<()> {
    let vec_fut = FuturesUnordered::new();

    for (svr_idx, svr_cfg) in context.config().server.iter().enumerate() {
        let context = context.clone();
        let flow_stat = flow_stat
            .get(svr_cfg.addr().port())
            .expect("port not existed in multi-server flow statistic")
            .clone();

        let svr_fut = listen(context, flow_stat, svr_idx);
        vec_fut.push(svr_fut);
    }

    match vec_fut.into_future().await.0 {
        Some(res) => {
            error!("one of UDP servers exited unexpectly, result: {:?}", res);
            let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
            Err(err)
        }
        None => unreachable!(),
    }
}
<commit_msg>Refine cache key for assoc_map<commit_after>//! UDP relay proxy server

use std::{
    io::{self, Cursor},
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use bytes::BytesMut;
use futures::{future, future::AbortHandle, stream::FuturesUnordered, StreamExt};
use log::{debug, error, info, trace, warn};
use lru_time_cache::{Entry, LruCache};
use tokio::{
    self,
    net::udp::{RecvHalf, SendHalf},
    sync::{mpsc, Mutex},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        flow::{SharedMultiServerFlowStatistic, SharedServerFlowStatistic},
        socks5::Address,
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

use super::{
    crypto_io::{decrypt_payload, encrypt_payload},
    DEFAULT_TIMEOUT,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

// Represent a UDP association
struct UdpAssociation {
    // local -> remote Queue
    // Drops tx, will close local -> remote task
    tx: mpsc::Sender<Vec<u8>>,

    // local <- remote task life watcher
    watcher: AbortHandle,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.watcher.abort();
    }
}

impl UdpAssociation {
    /// Create an association with addr
    async fn associate(
        context: SharedContext,
        svr_idx: usize,
        src_addr: SocketAddr,
        mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
    ) -> io::Result<UdpAssociation> {
        // Create a socket for receiving packets
        let local_addr = match context.config().local_addr {
            None => {
                // Let system allocate an address for us
                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
            }
            Some(ref addr) => {
                // Uses configured local address
                addr.bind_addr(&context).await?
            }
        };
        let remote_udp = create_udp_socket(&local_addr).await?;

        let local_addr = remote_udp.local_addr().expect("could not determine port bound to");
        debug!("created UDP Association for {} from {}", src_addr, local_addr);

        // Create a channel for sending packets to remote
        // FIXME: Channel size 1024?
        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(1024);

        // Splits socket into sender and receiver
        let (mut receiver, mut sender) = remote_udp.split();

        let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);

        // local -> remote
        {
            let context = context.clone();
            tokio::spawn(async move {
                let svr_cfg = context.server_config(svr_idx);

                while let Some(pkt) = rx.recv().await {
                    // pkt is already a raw packet, so just send it
                    if let Err(err) =
                        UdpAssociation::relay_l2r(&context, src_addr, &mut sender, &pkt[..], timeout, svr_cfg).await
                    {
                        error!("failed to relay packet, {} -> ..., error: {}", src_addr, err);

                        // FIXME: Ignore? Or how to deal with it?
                    }
                }

                debug!("UDP ASSOCIATE {} -> .. finished", src_addr);
            });
        }

        let (r2l_task, close_flag) = future::abortable(async move {
            let svr_cfg = context.server_config(svr_idx);

            loop {
                // Read and send back to source
                match UdpAssociation::relay_r2l(&context, src_addr, &mut receiver, &mut response_tx, svr_cfg).await {
                    Ok(..) => {}
                    Err(err) => {
                        error!("failed to receive packet, {} <- .., error: {}", src_addr, err);

                        // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                        // break;
                    }
                }
            }
        });

        // local <- remote
        tokio::spawn(async move {
            let _ = r2l_task.await;

            debug!("UDP ASSOCIATE {} <- .. finished", src_addr);
        });

        Ok(UdpAssociation {
            tx,
            watcher: close_flag,
        })
    }

    /// Relay packets from local to remote
    async fn relay_l2r(
        context: &Context,
        src: SocketAddr,
        remote_udp: &mut SendHalf,
        pkt: &[u8],
        timeout: Duration,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // First of all, decrypt payload CLIENT -> SERVER
        let decrypted_pkt = match decrypt_payload(context, svr_cfg.method(), svr_cfg.key(), pkt) {
            Ok(Some(pkt)) => pkt,
            Ok(None) => {
                error!("failed to decrypt pkt in UDP relay, packet too short");
                let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
                return Err(err);
            }
            Err(err) => {
                error!("failed to decrypt pkt in UDP relay: {}", err);
                let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed");
                return Err(err);
            }
        };

        // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
        let mut cur = Cursor::new(decrypted_pkt);

        let addr = Address::read_from(&mut cur).await?;

        debug!("UDP ASSOCIATE {} <-> {} establishing", src, addr);

        if context.check_outbound_blocked(&addr) {
            warn!("outbound {} is blocked by ACL rules", addr);
            return Ok(());
        }

        // Take out internal buffer for optimizing one byte copy
        let header_len = cur.position() as usize;
        let decrypted_pkt = cur.into_inner();
        let body = &decrypted_pkt[header_len..];

        let send_len = match addr {
            Address::SocketAddress(ref remote_addr) => {
                debug!(
                    "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                    src,
                    addr,
                    remote_addr,
                    body.len()
                );
                try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await?
            }
            Address::DomainNameAddress(ref dname, port) => lookup_outbound_then!(context, dname, port, |remote_addr| {
                match try_timeout(remote_udp.send_to(body, &remote_addr), Some(timeout)).await {
                    Ok(l) => {
                        debug!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Ok(l)
                    }
                    Err(err) => {
                        error!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Err(err)
                    }
                }
            })
            .map(|(_, l)| l)?,
        };

        assert_eq!(body.len(), send_len);

        Ok(())
    }

    /// Relay packets from remote to local
    async fn relay_r2l(
        context: &Context,
        src_addr: SocketAddr,
        remote_udp: &mut RecvHalf,
        response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // Waiting for response from server SERVER -> CLIENT
        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let mut remote_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        let (remote_recv_len, remote_addr) = remote_udp.recv_from(&mut remote_buf).await?;

        debug!(
            "UDP ASSOCIATE {} <- {}, payload length {} bytes",
            src_addr, remote_addr, remote_recv_len
        );

        // FIXME: The Address should be the Address that client sent
        let addr = Address::SocketAddress(remote_addr);

        // CLIENT <- SERVER protocol: ADDRESS + PAYLOAD
        let mut send_buf = Vec::new();
        addr.write_to_buf(&mut send_buf);
        send_buf.extend_from_slice(&remote_buf[..remote_recv_len]);

        let mut encrypt_buf = BytesMut::new();
        encrypt_payload(context, svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?;

        // Send back to src_addr
        if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await {
            error!("failed to send packet into response channel, error: {}", err);

            // FIXME: What to do? Ignore?
        }

        Ok(())
    }

    // Send packet to remote
    //
    // Return `Err` if receiver have been closed
    async fn send(&mut self, pkt: Vec<u8>) {
        if let Err(..) = self.tx.send(pkt).await {
            // SHOULDn't HAPPEN
            unreachable!("UDP Association local -> remote Queue closed unexpectly");
        }
    }
}

fn serialize(saddr: &SocketAddr) -> [u8; 18] {
    let mut result = [0; 18];
    result[..16].copy_from_slice(&match saddr.ip() {
        IpAddr::V4(ref ip) => ip.to_ipv6_mapped(),
        IpAddr::V6(ref ip) => ip.clone(),
    }.octets());
    result[16..].copy_from_slice(&saddr.port().to_ne_bytes());
    result
}

async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, svr_idx: usize) -> io::Result<()> {
    let svr_cfg = context.server_config(svr_idx);
    let listen_addr = svr_cfg.addr().bind_addr(&context).await?;

    let listener = create_udp_socket(&listen_addr).await?;
    let local_addr = listener.local_addr().expect("determine port bound to");
    info!("shadowsocks UDP listening on {}", local_addr);

    let (mut r, mut w) = listener.split();

    // NOTE: Associations are only eliminated by expire time by default
    // So it may exhaust all available file descriptors
    let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);
    let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations {
        LruCache::with_expiry_duration_and_capacity(timeout, max_assoc)
    } else {
        LruCache::with_expiry_duration(timeout)
    };
    let assoc_map = Arc::new(Mutex::new(assoc_map));

    // FIXME: Channel size 1024?
    let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024);

    {
        // Tokio task for sending data back to clients

        let assoc_map = assoc_map.clone();
        let flow_stat = flow_stat.clone();

        tokio::spawn(async move {
            while let Some((src, pkt)) = rx.recv().await {
                let cache_key = serialize(&src);
                {
                    let mut amap = assoc_map.lock().await;

                    // Check or update expire time
                    if amap.get(&cache_key).is_none() {
                        debug!(
                            "UDP association {} <-> ... is already expired, throwing away packet {} bytes",
                            src,
                            pkt.len()
                        );
                        continue;
                    }
                }

                if let Err(err) = w.send_to(&pkt, &src).await {
                    error!("UDP packet send failed, err: {:?}", err);
                    break;
                }

                flow_stat.udp().incr_tx(pkt.len() as u64);
            }

            // FIXME: How to stop the outer listener Future?
        });
    }

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src) = r.recv_from(&mut pkt_buf).await?;

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let pkt = &pkt_buf[..recv_len];

        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
        flow_stat.udp().incr_rx(pkt.len() as u64);

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check ACL
        if context.check_client_blocked(&src) {
            warn!("client {} is blocked by ACL rules", src);
            continue;
        }

        // Check or (re)create an association
        {
            // Locks the whole association map
            let mut assoc_map = assoc_map.lock().await;

            // Get or create an association
            let assoc = match assoc_map.entry(serialize(&src)) {
                Entry::Occupied(oc) => oc.into_mut(),
                Entry::Vacant(vc) => vc.insert(
                    UdpAssociation::associate(context.clone(), svr_idx, src, tx.clone())
                        .await
                        .expect("create udp association"),
                ),
            };

            // FIXME: Lock is still kept for a mutable reference
            // Send to local -> remote task
            assoc.send(pkt.to_vec()).await;
        }
    }
}

/// Starts a UDP relay server
pub async fn run(context: SharedContext, flow_stat: SharedMultiServerFlowStatistic) -> io::Result<()> {
    let vec_fut = FuturesUnordered::new();

    for (svr_idx, svr_cfg) in context.config().server.iter().enumerate() {
        let context = context.clone();
        let flow_stat = flow_stat
            .get(svr_cfg.addr().port())
            .expect("port not existed in multi-server flow statistic")
            .clone();

        let svr_fut = listen(context, flow_stat, svr_idx);
        vec_fut.push(svr_fut);
    }

    match vec_fut.into_future().await.0 {
        Some(res) => {
            error!("one of UDP servers exited unexpectly, result: {:?}", res);
            let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
            Err(err)
        }
        None => unreachable!(),
    }
}
<|endoftext|>


<commit_before>//! UDP relay proxy server

use std::{
    io::{self, Cursor},
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use bytes::BytesMut;
use futures::{future, future::AbortHandle, stream::FuturesUnordered, StreamExt};
use log::{debug, error, info, trace, warn};
use lru_time_cache::{Entry, LruCache};
use tokio::{
    self,
    net::udp::{RecvHalf, SendHalf},
    sync::{mpsc, Mutex},
    time,
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        flow::{SharedMultiServerFlowStatistic, SharedServerFlowStatistic},
        socks5::Address,
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

use super::{
    crypto_io::{decrypt_payload, encrypt_payload},
    DEFAULT_TIMEOUT,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

// Represent a UDP association
struct UdpAssociation {
    // local -> remote Queue
    // Drops tx, will close local -> remote task
    tx: mpsc::Sender<Vec<u8>>,

    // local <- remote task life watcher
    watcher: AbortHandle,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.watcher.abort();
    }
}

impl UdpAssociation {
    /// Create an association with addr
    async fn associate(
        context: SharedContext,
        svr_idx: usize,
        src_addr: SocketAddr,
        mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
    ) -> io::Result<UdpAssociation> {
        // Create a socket for receiving packets
        let local_addr = match context.config().local_addr {
            None => {
                // Let system allocate an address for us
                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
            }
            Some(ref addr) => {
                // Uses configured local address
                addr.bind_addr(&context).await?
            }
        };
        let remote_udp = create_udp_socket(&local_addr).await?;

        let local_addr = remote_udp.local_addr().expect("could not determine port bound to");
        debug!("created UDP Association for {} from {}", src_addr, local_addr);

        // Create a channel for sending packets to remote
        // FIXME: Channel size 1024?
        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(1024);

        // Splits socket into sender and receiver
        let (mut receiver, mut sender) = remote_udp.split();

        let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);

        // local -> remote
        {
            let context = context.clone();
            tokio::spawn(async move {
                let svr_cfg = context.server_config(svr_idx);

                while let Some(pkt) = rx.recv().await {
                    // pkt is already a raw packet, so just send it
                    if let Err(err) =
                        UdpAssociation::relay_l2r(&context, src_addr, &mut sender, &pkt[..], timeout, svr_cfg).await
                    {
                        error!("failed to relay packet, {} -> ..., error: {}", src_addr, err);

                        // FIXME: Ignore? Or how to deal with it?
                    }
                }

                debug!("UDP ASSOCIATE {} -> .. finished", src_addr);
            });
        }

        let (r2l_task, close_flag) = future::abortable(async move {
            let svr_cfg = context.server_config(svr_idx);

            loop {
                // Read and send back to source
                match UdpAssociation::relay_r2l(&context, src_addr, &mut receiver, &mut response_tx, svr_cfg).await {
                    Ok(..) => {}
                    Err(err) => {
                        error!("failed to receive packet, {} <- .., error: {}", src_addr, err);

                        // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                        // break;
                    }
                }
            }
        });

        // local <- remote
        tokio::spawn(async move {
            let _ = r2l_task.await;

            debug!("UDP ASSOCIATE {} <- .. finished", src_addr);
        });

        Ok(UdpAssociation {
            tx,
            watcher: close_flag,
        })
    }

    /// Relay packets from local to remote
    async fn relay_l2r(
        context: &Context,
        src: SocketAddr,
        remote_udp: &mut SendHalf,
        pkt: &[u8],
        timeout: Duration,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // First of all, decrypt payload CLIENT -> SERVER
        let decrypted_pkt = match decrypt_payload(context, svr_cfg.method(), svr_cfg.key(), pkt) {
            Ok(Some(pkt)) => pkt,
            Ok(None) => {
                error!("failed to decrypt pkt in UDP relay, packet too short");
                let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
                return Err(err);
            }
            Err(err) => {
                error!("failed to decrypt pkt in UDP relay: {}", err);
                let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed");
                return Err(err);
            }
        };

        // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
        let mut cur = Cursor::new(decrypted_pkt);

        let addr = Address::read_from(&mut cur).await?;

        debug!("UDP ASSOCIATE {} <-> {} establishing", src, addr);

        if context.check_outbound_blocked(&addr) {
            warn!("outbound {} is blocked by ACL rules", addr);
            return Ok(());
        }

        // Take out internal buffer for optimizing one byte copy
        let header_len = cur.position() as usize;
        let decrypted_pkt = cur.into_inner();
        let body = &decrypted_pkt[header_len..];

        let send_len = match addr {
            Address::SocketAddress(ref remote_addr) => {
                debug!(
                    "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                    src,
                    addr,
                    remote_addr,
                    body.len()
                );
                try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await?
            }
            Address::DomainNameAddress(ref dname, port) => lookup_outbound_then!(context, dname, port, |remote_addr| {
                match try_timeout(remote_udp.send_to(body, &remote_addr), Some(timeout)).await {
                    Ok(l) => {
                        debug!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Ok(l)
                    }
                    Err(err) => {
                        error!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Err(err)
                    }
                }
            })
            .map(|(_, l)| l)?,
        };

        assert_eq!(body.len(), send_len);

        Ok(())
    }

    /// Relay packets from remote to local
    async fn relay_r2l(
        context: &Context,
        src_addr: SocketAddr,
        remote_udp: &mut RecvHalf,
        response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // Waiting for response from server SERVER -> CLIENT
        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let mut remote_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        let (remote_recv_len, remote_addr) = remote_udp.recv_from(&mut remote_buf).await?;

        debug!(
            "UDP ASSOCIATE {} <- {}, payload length {} bytes",
            src_addr, remote_addr, remote_recv_len
        );

        // FIXME: The Address should be the Address that client sent
        let addr = Address::SocketAddress(remote_addr);

        // CLIENT <- SERVER protocol: ADDRESS + PAYLOAD
        let mut send_buf = Vec::new();
        addr.write_to_buf(&mut send_buf);
        send_buf.extend_from_slice(&remote_buf[..remote_recv_len]);

        let mut encrypt_buf = BytesMut::new();
        encrypt_payload(context, svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?;

        // Send back to src_addr
        if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await {
            error!("failed to send packet into response channel, error: {}", err);

            // FIXME: What to do? Ignore?
        }

        Ok(())
    }

    // Send packet to remote
    //
    // Return `Err` if receiver have been closed
    async fn send(&mut self, pkt: Vec<u8>) {
        if let Err(..) = self.tx.send(pkt).await {
            // SHOULDn't HAPPEN
            unreachable!("UDP Association local -> remote Queue closed unexpectly");
        }
    }
}

async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, svr_idx: usize) -> io::Result<()> {
    let svr_cfg = context.server_config(svr_idx);
    let listen_addr = svr_cfg.addr().bind_addr(&context).await?;

    let listener = create_udp_socket(&listen_addr).await?;
    let local_addr = listener.local_addr().expect("determine port bound to");
    info!("shadowsocks UDP listening on {}", local_addr);

    let (mut r, mut w) = listener.split();

    // NOTE: Associations are only eliminated by expire time by default
    // So it may exhaust all available file descriptors
    let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);
    let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations {
        LruCache::with_expiry_duration_and_capacity(timeout, max_assoc)
    } else {
        LruCache::with_expiry_duration(timeout)
    };
    let assoc_map = Arc::new(Mutex::new(assoc_map));

    // FIXME: Channel size 1024?
    let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024);

    {
        // Tokio task for sending data back to clients

        let assoc_map = assoc_map.clone();
        let flow_stat = flow_stat.clone();

        tokio::spawn(async move {
            while let Some((src, pkt)) = rx.recv().await {
                let cache_key = src.to_string();
                {
                    let mut amap = assoc_map.lock().await;

                    // Check or update expire time
                    if amap.get(&cache_key).is_none() {
                        debug!(
                            "UDP association {} <-> ... is already expired, throwing away packet {} bytes",
                            src,
                            pkt.len()
                        );
                        continue;
                    }
                }

                if let Err(err) = w.send_to(&pkt, &src).await {
                    error!("UDP packet send failed, err: {:?}", err);
                    break;
                }

                flow_stat.udp().incr_tx(pkt.len() as u64);
            }

            // FIXME: How to stop the outer listener Future?
        });
    }

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src) = match time::timeout(timeout, r.recv_from(&mut pkt_buf)).await {
            Ok(r) => r?,
            Err(..) => {
                // Cleanup expired association
                // Do not consume this iterator, it will updates expire time of items that traversed
                let mut assoc_map = assoc_map.lock().await;
                let _ = assoc_map.iter();
                continue;
            }
        };

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let pkt = &pkt_buf[..recv_len];

        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
        flow_stat.udp().incr_rx(pkt.len() as u64);

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check ACL
        if context.check_client_blocked(&src) {
            warn!("client {} is blocked by ACL rules", src);
            continue;
        }

        // Check or (re)create an association
        {
            // Locks the whole association map
            let mut assoc_map = assoc_map.lock().await;

            // Get or create an association
            let assoc = match assoc_map.entry(src.to_string()) {
                Entry::Occupied(oc) => oc.into_mut(),
                Entry::Vacant(vc) => vc.insert(
                    UdpAssociation::associate(context.clone(), svr_idx, src, tx.clone())
                        .await
                        .expect("create udp association"),
                ),
            };

            // FIXME: Lock is still kept for a mutable reference
            // Send to local -> remote task
            assoc.send(pkt.to_vec()).await;
        }
    }
}

/// Starts a UDP relay server
pub async fn run(context: SharedContext, flow_stat: SharedMultiServerFlowStatistic) -> io::Result<()> {
    let vec_fut = FuturesUnordered::new();

    for (svr_idx, svr_cfg) in context.config().server.iter().enumerate() {
        let context = context.clone();
        let flow_stat = flow_stat
            .get(svr_cfg.addr().port())
            .expect("port not existed in multi-server flow statistic")
            .clone();

        let svr_fut = listen(context, flow_stat, svr_idx);
        vec_fut.push(svr_fut);
    }

    match vec_fut.into_future().await.0 {
        Some(res) => {
            error!("one of UDP servers exited unexpectly, result: {:?}", res);
            let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
            Err(err)
        }
        None => unreachable!(),
    }
}
<commit_msg>Remove unnecessary wakeups from UDP<commit_after>//! UDP relay proxy server

use std::{
    io::{self, Cursor},
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use bytes::BytesMut;
use futures::{future, future::AbortHandle, stream::FuturesUnordered, StreamExt};
use log::{debug, error, info, trace, warn};
use lru_time_cache::{Entry, LruCache};
use tokio::{
    self,
    net::udp::{RecvHalf, SendHalf},
    sync::{mpsc, Mutex},
};

use crate::{
    config::ServerConfig,
    context::{Context, SharedContext},
    relay::{
        flow::{SharedMultiServerFlowStatistic, SharedServerFlowStatistic},
        socks5::Address,
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

use super::{
    crypto_io::{decrypt_payload, encrypt_payload},
    DEFAULT_TIMEOUT,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

// Represent a UDP association
struct UdpAssociation {
    // local -> remote Queue
    // Drops tx, will close local -> remote task
    tx: mpsc::Sender<Vec<u8>>,

    // local <- remote task life watcher
    watcher: AbortHandle,
}

impl Drop for UdpAssociation {
    fn drop(&mut self) {
        self.watcher.abort();
    }
}

impl UdpAssociation {
    /// Create an association with addr
    async fn associate(
        context: SharedContext,
        svr_idx: usize,
        src_addr: SocketAddr,
        mut response_tx: mpsc::Sender<(SocketAddr, BytesMut)>,
    ) -> io::Result<UdpAssociation> {
        // Create a socket for receiving packets
        let local_addr = match context.config().local_addr {
            None => {
                // Let system allocate an address for us
                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
            }
            Some(ref addr) => {
                // Uses configured local address
                addr.bind_addr(&context).await?
            }
        };
        let remote_udp = create_udp_socket(&local_addr).await?;

        let local_addr = remote_udp.local_addr().expect("could not determine port bound to");
        debug!("created UDP Association for {} from {}", src_addr, local_addr);

        // Create a channel for sending packets to remote
        // FIXME: Channel size 1024?
        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(1024);

        // Splits socket into sender and receiver
        let (mut receiver, mut sender) = remote_udp.split();

        let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);

        // local -> remote
        {
            let context = context.clone();
            tokio::spawn(async move {
                let svr_cfg = context.server_config(svr_idx);

                while let Some(pkt) = rx.recv().await {
                    // pkt is already a raw packet, so just send it
                    if let Err(err) =
                        UdpAssociation::relay_l2r(&context, src_addr, &mut sender, &pkt[..], timeout, svr_cfg).await
                    {
                        error!("failed to relay packet, {} -> ..., error: {}", src_addr, err);

                        // FIXME: Ignore? Or how to deal with it?
                    }
                }

                debug!("UDP ASSOCIATE {} -> .. finished", src_addr);
            });
        }

        let (r2l_task, close_flag) = future::abortable(async move {
            let svr_cfg = context.server_config(svr_idx);

            loop {
                // Read and send back to source
                match UdpAssociation::relay_r2l(&context, src_addr, &mut receiver, &mut response_tx, svr_cfg).await {
                    Ok(..) => {}
                    Err(err) => {
                        error!("failed to receive packet, {} <- .., error: {}", src_addr, err);

                        // FIXME: Don't break, or if you can find a way to drop the UdpAssociation
                        // break;
                    }
                }
            }
        });

        // local <- remote
        tokio::spawn(async move {
            let _ = r2l_task.await;

            debug!("UDP ASSOCIATE {} <- .. finished", src_addr);
        });

        Ok(UdpAssociation {
            tx,
            watcher: close_flag,
        })
    }

    /// Relay packets from local to remote
    async fn relay_l2r(
        context: &Context,
        src: SocketAddr,
        remote_udp: &mut SendHalf,
        pkt: &[u8],
        timeout: Duration,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // First of all, decrypt payload CLIENT -> SERVER
        let decrypted_pkt = match decrypt_payload(context, svr_cfg.method(), svr_cfg.key(), pkt) {
            Ok(Some(pkt)) => pkt,
            Ok(None) => {
                error!("failed to decrypt pkt in UDP relay, packet too short");
                let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
                return Err(err);
            }
            Err(err) => {
                error!("failed to decrypt pkt in UDP relay: {}", err);
                let err = io::Error::new(io::ErrorKind::InvalidData, "decrypt failed");
                return Err(err);
            }
        };

        // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
        let mut cur = Cursor::new(decrypted_pkt);

        let addr = Address::read_from(&mut cur).await?;

        debug!("UDP ASSOCIATE {} <-> {} establishing", src, addr);

        if context.check_outbound_blocked(&addr) {
            warn!("outbound {} is blocked by ACL rules", addr);
            return Ok(());
        }

        // Take out internal buffer for optimizing one byte copy
        let header_len = cur.position() as usize;
        let decrypted_pkt = cur.into_inner();
        let body = &decrypted_pkt[header_len..];

        let send_len = match addr {
            Address::SocketAddress(ref remote_addr) => {
                debug!(
                    "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                    src,
                    addr,
                    remote_addr,
                    body.len()
                );
                try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await?
            }
            Address::DomainNameAddress(ref dname, port) => lookup_outbound_then!(context, dname, port, |remote_addr| {
                match try_timeout(remote_udp.send_to(body, &remote_addr), Some(timeout)).await {
                    Ok(l) => {
                        debug!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Ok(l)
                    }
                    Err(err) => {
                        error!(
                            "UDP ASSOCIATE {} -> {} ({}), payload length {} bytes",
                            src,
                            addr,
                            remote_addr,
                            body.len()
                        );
                        Err(err)
                    }
                }
            })
            .map(|(_, l)| l)?,
        };

        assert_eq!(body.len(), send_len);

        Ok(())
    }

    /// Relay packets from remote to local
    async fn relay_r2l(
        context: &Context,
        src_addr: SocketAddr,
        remote_udp: &mut RecvHalf,
        response_tx: &mut mpsc::Sender<(SocketAddr, BytesMut)>,
        svr_cfg: &ServerConfig,
    ) -> io::Result<()> {
        // Waiting for response from server SERVER -> CLIENT
        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let mut remote_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
        let (remote_recv_len, remote_addr) = remote_udp.recv_from(&mut remote_buf).await?;

        debug!(
            "UDP ASSOCIATE {} <- {}, payload length {} bytes",
            src_addr, remote_addr, remote_recv_len
        );

        // FIXME: The Address should be the Address that client sent
        let addr = Address::SocketAddress(remote_addr);

        // CLIENT <- SERVER protocol: ADDRESS + PAYLOAD
        let mut send_buf = Vec::new();
        addr.write_to_buf(&mut send_buf);
        send_buf.extend_from_slice(&remote_buf[..remote_recv_len]);

        let mut encrypt_buf = BytesMut::new();
        encrypt_payload(context, svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?;

        // Send back to src_addr
        if let Err(err) = response_tx.send((src_addr, encrypt_buf)).await {
            error!("failed to send packet into response channel, error: {}", err);

            // FIXME: What to do? Ignore?
        }

        Ok(())
    }

    // Send packet to remote
    //
    // Return `Err` if receiver have been closed
    async fn send(&mut self, pkt: Vec<u8>) {
        if let Err(..) = self.tx.send(pkt).await {
            // SHOULDn't HAPPEN
            unreachable!("UDP Association local -> remote Queue closed unexpectly");
        }
    }
}

async fn listen(context: SharedContext, flow_stat: SharedServerFlowStatistic, svr_idx: usize) -> io::Result<()> {
    let svr_cfg = context.server_config(svr_idx);
    let listen_addr = svr_cfg.addr().bind_addr(&context).await?;

    let listener = create_udp_socket(&listen_addr).await?;
    let local_addr = listener.local_addr().expect("determine port bound to");
    info!("shadowsocks UDP listening on {}", local_addr);

    let (mut r, mut w) = listener.split();

    // NOTE: Associations are only eliminated by expire time by default
    // So it may exhaust all available file descriptors
    let timeout = context.config().udp_timeout.unwrap_or(DEFAULT_TIMEOUT);
    let assoc_map = if let Some(max_assoc) = context.config().udp_max_associations {
        LruCache::with_expiry_duration_and_capacity(timeout, max_assoc)
    } else {
        LruCache::with_expiry_duration(timeout)
    };
    let assoc_map = Arc::new(Mutex::new(assoc_map));

    // FIXME: Channel size 1024?
    let (tx, mut rx) = mpsc::channel::<(SocketAddr, BytesMut)>(1024);

    {
        // Tokio task for sending data back to clients

        let assoc_map = assoc_map.clone();
        let flow_stat = flow_stat.clone();

        tokio::spawn(async move {
            while let Some((src, pkt)) = rx.recv().await {
                let cache_key = src.to_string();
                {
                    let mut amap = assoc_map.lock().await;

                    // Check or update expire time
                    if amap.get(&cache_key).is_none() {
                        debug!(
                            "UDP association {} <-> ... is already expired, throwing away packet {} bytes",
                            src,
                            pkt.len()
                        );
                        continue;
                    }
                }

                if let Err(err) = w.send_to(&pkt, &src).await {
                    error!("UDP packet send failed, err: {:?}", err);
                    break;
                }

                flow_stat.udp().incr_tx(pkt.len() as u64);
            }

            // FIXME: How to stop the outer listener Future?
        });
    }

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src) = r.recv_from(&mut pkt_buf).await?;

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        let pkt = &pkt_buf[..recv_len];

        trace!("received UDP packet from {}, length {} bytes", src, recv_len);
        flow_stat.udp().incr_rx(pkt.len() as u64);

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check ACL
        if context.check_client_blocked(&src) {
            warn!("client {} is blocked by ACL rules", src);
            continue;
        }

        // Check or (re)create an association
        {
            // Locks the whole association map
            let mut assoc_map = assoc_map.lock().await;

            // Get or create an association
            let assoc = match assoc_map.entry(src.to_string()) {
                Entry::Occupied(oc) => oc.into_mut(),
                Entry::Vacant(vc) => vc.insert(
                    UdpAssociation::associate(context.clone(), svr_idx, src, tx.clone())
                        .await
                        .expect("create udp association"),
                ),
            };

            // FIXME: Lock is still kept for a mutable reference
            // Send to local -> remote task
            assoc.send(pkt.to_vec()).await;
        }
    }
}

/// Starts a UDP relay server
pub async fn run(context: SharedContext, flow_stat: SharedMultiServerFlowStatistic) -> io::Result<()> {
    let vec_fut = FuturesUnordered::new();

    for (svr_idx, svr_cfg) in context.config().server.iter().enumerate() {
        let context = context.clone();
        let flow_stat = flow_stat
            .get(svr_cfg.addr().port())
            .expect("port not existed in multi-server flow statistic")
            .clone();

        let svr_fut = listen(context, flow_stat, svr_idx);
        vec_fut.push(svr_fut);
    }

    match vec_fut.into_future().await.0 {
        Some(res) => {
            error!("one of UDP servers exited unexpectly, result: {:?}", res);
            let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
            Err(err)
        }
        None => unreachable!(),
    }
}
<|endoftext|>