package utp import ( "context" "errors" "io" "log" "math/rand" "net" "sync" "time" "github.com/anacrolix/missinggo" "github.com/anacrolix/missinggo/inproc" "github.com/anacrolix/missinggo/pproffd" ) var ( _ net.Listener = &Socket{} _ net.PacketConn = &Socket{} ) // Uniquely identifies any uTP connection on top of the underlying packet // stream. type connKey struct { remoteAddr resolvedAddrStr connID uint16 } // A Socket wraps a net.PacketConn, diverting uTP packets to its child uTP // Conns. type Socket struct { pc net.PacketConn conns map[connKey]*Conn backlogNotEmpty missinggo.Event backlog map[syn]struct{} closed missinggo.Event destroyed missinggo.Event wgReadWrite sync.WaitGroup unusedReads chan read connDeadlines // If a read error occurs on the underlying net.PacketConn, it is put // here. This is because reading is done in its own goroutine to dispatch // to uTP Conns. ReadErr error } func listenPacket(network, addr string) (pc net.PacketConn, err error) { if network == "inproc" { return inproc.ListenPacket(network, addr) } return net.ListenPacket(network, addr) } // NewSocket creates a net.PacketConn with the given network and address, and // returns a Socket dispatching on it. func NewSocket(network, addr string) (s *Socket, err error) { if network == "" { network = "udp" } pc, err := listenPacket(network, addr) if err != nil { return } return NewSocketFromPacketConn(pc) } // Create a Socket, using the provided net.PacketConn. If you want to retain // use of the net.PacketConn after the Socket closes it, override the // net.PacketConn's Close method, or use NetSocketFromPacketConnNoClose. func NewSocketFromPacketConn(pc net.PacketConn) (s *Socket, err error) { s = &Socket{ backlog: make(map[syn]struct{}, backlog), pc: pc, unusedReads: make(chan read, 100), wgReadWrite: sync.WaitGroup{}, } mu.Lock() sockets[s] = struct{}{} mu.Unlock() go s.reader() return } // Create a Socket using the provided PacketConn, that doesn't close the // PacketConn when the Socket is closed. func NewSocketFromPacketConnNoClose(pc net.PacketConn) (s *Socket, err error) { return NewSocketFromPacketConn(packetConnNopCloser{pc}) } func (s *Socket) unusedRead(read read) { unusedReads.Add(1) select { case s.unusedReads <- read: default: // Drop the packet. unusedReadsDropped.Add(1) } } func (s *Socket) strNetAddr(str string) (a net.Addr) { var err error switch n := s.network(); n { case "udp": a, err = net.ResolveUDPAddr(n, str) case "inproc": a, err = inproc.ResolveAddr(n, str) default: panic(n) } if err != nil { panic(err) } return } func (s *Socket) pushBacklog(syn syn) { if _, ok := s.backlog[syn]; ok { return } // Pop a pseudo-random syn to make room. TODO: Use missinggo/orderedmap, // coz that's what is wanted here. for k := range s.backlog { if len(s.backlog) < backlog { break } delete(s.backlog, k) // A syn is sent on the remote's recv_id, so this is where we can send // the reset. s.reset(s.strNetAddr(k.addr), k.seq_nr, k.conn_id) } s.backlog[syn] = struct{}{} s.backlogChanged() } func (s *Socket) reader() { mu.Lock() defer mu.Unlock() defer s.destroy() var b [maxRecvSize]byte for { s.wgReadWrite.Add(1) mu.Unlock() n, addr, err := s.pc.ReadFrom(b[:]) s.wgReadWrite.Done() mu.Lock() if s.destroyed.IsSet() { return } if err != nil { log.Printf("error reading Socket PacketConn: %s", err) s.ReadErr = err return } s.handleReceivedPacket(read{ append([]byte(nil), b[:n]...), addr, }) } } func receivedUTPPacketSize(n int) { if n > largestReceivedUTPPacket { largestReceivedUTPPacket = n largestReceivedUTPPacketExpvar.Set(int64(n)) } } func (s *Socket) connForRead(h header, from net.Addr) (c *Conn, ok bool) { c, ok = s.conns[connKey{ resolvedAddrStr(from.String()), func() uint16 { if h.Type == stSyn { // SYNs have a ConnID one lower than the eventual recvID, and we index // the connections with that, so use it for the lookup. return h.ConnID + 1 } else { return h.ConnID } }(), }] return } func (s *Socket) handlePacketReceivedForEstablishedConn(h header, from net.Addr, data []byte, c *Conn) { if h.Type == stSyn { if h.ConnID == c.send_id-2 { // This is a SYN for connection that cannot exist locally. The // connection the remote wants to establish here with the proposed // recv_id, already has an existing connection that was dialled // *out* from this socket, which is why the send_id is 1 higher, // rather than 1 lower than the recv_id. log.Print("resetting conflicting syn") s.reset(from, h.SeqNr, h.ConnID) return } else if h.ConnID != c.send_id { panic("bad assumption") } } c.receivePacket(h, data) } func (s *Socket) handleReceivedPacket(p read) { if len(p.data) < 20 { s.unusedRead(p) return } var h header hEnd, err := h.Unmarshal(p.data) if err != nil || h.Type > stMax || h.Version != 1 { s.unusedRead(p) return } if c, ok := s.connForRead(h, p.from); ok { receivedUTPPacketSize(len(p.data)) s.handlePacketReceivedForEstablishedConn(h, p.from, p.data[hEnd:], c) return } // Packet doesn't belong to an existing connection. switch h.Type { case stSyn: s.pushBacklog(syn{ seq_nr: h.SeqNr, conn_id: h.ConnID, addr: p.from.String(), }) return case stReset: // Could be a late arriving packet for a Conn we're already done with. // If it was for an existing connection, we would have handled it // earlier. default: unexpectedPacketsRead.Add(1) // This is an unexpected packet. We'll send a reset, but also pass it // on. I don't think you can reset on the received packets ConnID if // it isn't a SYN, as the send_id will differ in this case. s.reset(p.from, h.SeqNr, h.ConnID) // Connection initiated by remote. s.reset(p.from, h.SeqNr, h.ConnID-1) // Connection initiated locally. s.reset(p.from, h.SeqNr, h.ConnID+1) } s.unusedRead(p) } // Send a reset in response to a packet with the given header. func (s *Socket) reset(addr net.Addr, ackNr, connId uint16) { b := make([]byte, 0, maxHeaderSize) h := header{ Type: stReset, Version: 1, ConnID: connId, AckNr: ackNr, } b = b[:h.Marshal(b)] go s.writeTo(b, addr) } // Return a recv_id that should be free. Handling the case where it isn't is // deferred to a more appropriate function. func (s *Socket) newConnID(remoteAddr resolvedAddrStr) (id uint16) { // Rather than use math.Rand, which requires generating all the IDs up // front and allocating a slice, we do it on the stack, generating the IDs // only as required. To do this, we use the fact that the array is // default-initialized. IDs that are 0, are actually their index in the // array. IDs that are non-zero, are +1 from their intended ID. var idsBack [0x10000]int ids := idsBack[:] for len(ids) != 0 { // Pick the next ID from the untried ids. i := rand.Intn(len(ids)) id = uint16(ids[i]) // If it's zero, then treat it as though the index i was the ID. // Otherwise the value we get is the ID+1. if id == 0 { id = uint16(i) } else { id-- } // Check there's no connection using this ID for its recv_id... _, ok1 := s.conns[connKey{remoteAddr, id}] // and if we're connecting to our own Socket, that there isn't a Conn // already receiving on what will correspond to our send_id. Note that // we just assume that we could be connecting to our own Socket. This // will halve the available connection IDs to each distinct remote // address. Presumably that's ~0x8000, down from ~0x10000. _, ok2 := s.conns[connKey{remoteAddr, id + 1}] _, ok4 := s.conns[connKey{remoteAddr, id - 1}] if !ok1 && !ok2 && !ok4 { return } // The set of possible IDs is shrinking. The highest one will be lost, so // it's moved to the location of the one we just tried. ids[i] = len(ids) // Conveniently already +1. // And shrink. ids = ids[:len(ids)-1] } return } var ( zeroipv4 = net.ParseIP("0.0.0.0") zeroipv6 = net.ParseIP("::") ipv4lo = mustResolveUDP("127.0.0.1") ipv6lo = mustResolveUDP("::1") ) func mustResolveUDP(addr string) net.IP { u, err := net.ResolveIPAddr("ip", addr) if err != nil { panic(err) } return u.IP } func realRemoteAddr(addr net.Addr) net.Addr { udpAddr, ok := addr.(*net.UDPAddr) if ok { if udpAddr.IP.Equal(zeroipv4) { udpAddr.IP = ipv4lo } if udpAddr.IP.Equal(zeroipv6) { udpAddr.IP = ipv6lo } } return addr } func (s *Socket) newConn(addr net.Addr) (c *Conn) { addr = realRemoteAddr(addr) c = &Conn{ socket: s, remoteSocketAddr: addr, created: time.Now(), } c.sendPendingSendSendStateTimer = missinggo.StoppedFuncTimer(c.sendPendingSendStateTimerCallback) c.packetReadTimeoutTimer = time.AfterFunc(packetReadTimeout, c.receivePacketTimeoutCallback) return } func (s *Socket) Dial(addr string) (net.Conn, error) { return s.DialContext(context.Background(), "", addr) } func (s *Socket) resolveAddr(network, addr string) (net.Addr, error) { n := s.network() if network != "" { n = network } if n == "inproc" { return inproc.ResolveAddr(n, addr) } return net.ResolveUDPAddr(n, addr) } func (s *Socket) network() string { return s.pc.LocalAddr().Network() } func (s *Socket) startOutboundConn(addr net.Addr) (c *Conn, err error) { mu.Lock() defer mu.Unlock() c = s.newConn(addr) c.recv_id = s.newConnID(resolvedAddrStr(c.RemoteAddr().String())) c.send_id = c.recv_id + 1 if logLevel >= 1 { log.Printf("dial registering addr: %s", c.RemoteAddr().String()) } if !s.registerConn(c.recv_id, resolvedAddrStr(c.RemoteAddr().String()), c) { err = errors.New("couldn't register new connection") log.Println(c.recv_id, c.RemoteAddr().String()) for k, c := range s.conns { log.Println(k, c, c.age()) } log.Printf("that's %d connections", len(s.conns)) } if err != nil { return } c.seq_nr = 1 c.writeSyn() return } func (s *Socket) DialContext(ctx context.Context, network, addr string) (nc net.Conn, err error) { netAddr, err := s.resolveAddr(network, addr) if err != nil { return } c, err := s.startOutboundConn(netAddr) if err != nil { return } connErr := make(chan error, 1) go func() { connErr <- c.recvSynAck() }() select { case err = <-connErr: case <-ctx.Done(): err = ctx.Err() } if err != nil { mu.Lock() c.destroy(errors.New("dial timeout")) mu.Unlock() return } mu.Lock() c.updateCanWrite() mu.Unlock() nc = pproffd.WrapNetConn(c) return } func (me *Socket) writeTo(b []byte, addr net.Addr) (n int, err error) { apdc := artificialPacketDropChance if apdc != 0 { if rand.Float64() < apdc { n = len(b) return } } n, err = me.pc.WriteTo(b, addr) return } // Returns true if the connection was newly registered, false otherwise. func (s *Socket) registerConn(recvID uint16, remoteAddr resolvedAddrStr, c *Conn) bool { if s.conns == nil { s.conns = make(map[connKey]*Conn) } key := connKey{remoteAddr, recvID} if _, ok := s.conns[key]; ok { return false } c.connKey = key s.conns[key] = c return true } func (s *Socket) backlogChanged() { if len(s.backlog) != 0 { s.backlogNotEmpty.Set() } else { s.backlogNotEmpty.Clear() } } func (s *Socket) nextSyn() (syn syn, err error) { for { missinggo.WaitEvents(&mu, &s.closed, &s.backlogNotEmpty, &s.destroyed) if s.closed.IsSet() { err = errClosed return } if s.destroyed.IsSet() { err = s.ReadErr return } for k := range s.backlog { syn = k delete(s.backlog, k) s.backlogChanged() return } } } // ACK a SYN, and return a new Conn for it. ok is false if the SYN is bad, and // the Conn invalid. func (s *Socket) ackSyn(syn syn) (c *Conn, ok bool) { c = s.newConn(s.strNetAddr(syn.addr)) c.send_id = syn.conn_id c.recv_id = c.send_id + 1 c.seq_nr = uint16(rand.Int()) c.lastAck = c.seq_nr - 1 c.ack_nr = syn.seq_nr c.synAcked = true c.updateCanWrite() if !s.registerConn(c.recv_id, resolvedAddrStr(syn.addr), c) { // SYN that triggered this accept duplicates existing connection. // Ack again in case the SYN was a resend. c = s.conns[connKey{resolvedAddrStr(syn.addr), c.recv_id}] if c.send_id != syn.conn_id { panic(":|") } c.sendState() return } c.sendState() ok = true return } // Accept and return a new uTP connection. func (s *Socket) Accept() (net.Conn, error) { mu.Lock() defer mu.Unlock() for { syn, err := s.nextSyn() if err != nil { return nil, err } c, ok := s.ackSyn(syn) if ok { c.updateCanWrite() return c, nil } } } // The address we're listening on for new uTP connections. func (s *Socket) Addr() net.Addr { return s.pc.LocalAddr() } func (s *Socket) CloseNow() error { mu.Lock() defer mu.Unlock() s.closed.Set() for _, c := range s.conns { c.closeNow() } s.destroy() s.wgReadWrite.Wait() return nil } func (s *Socket) Close() error { mu.Lock() defer mu.Unlock() s.closed.Set() s.lazyDestroy() return nil } func (s *Socket) lazyDestroy() { if len(s.conns) != 0 { return } if !s.closed.IsSet() { return } s.destroy() } func (s *Socket) destroy() { delete(sockets, s) s.destroyed.Set() s.pc.Close() for _, c := range s.conns { c.destroy(errors.New("Socket destroyed")) } } func (s *Socket) LocalAddr() net.Addr { return s.pc.LocalAddr() } func (s *Socket) ReadFrom(p []byte) (n int, addr net.Addr, err error) { select { case read, ok := <-s.unusedReads: if !ok { err = io.EOF return } n = copy(p, read.data) addr = read.from return case <-s.connDeadlines.read.passed.LockedChan(&mu): err = errTimeout return } } func (s *Socket) WriteTo(b []byte, addr net.Addr) (n int, err error) { mu.Lock() if s.connDeadlines.write.passed.IsSet() { err = errTimeout } s.wgReadWrite.Add(1) defer s.wgReadWrite.Done() mu.Unlock() if err != nil { return } return s.pc.WriteTo(b, addr) }