gtsocial-umbx

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README | LICENSE

udp.go (2985B)

      1 //go:build !windows
      2 // +build !windows
      3 
      4 package dns
      5 
      6 import (
      7 	"net"
      8 
      9 	"golang.org/x/net/ipv4"
     10 	"golang.org/x/net/ipv6"
     11 )
     12 
     13 // This is the required size of the OOB buffer to pass to ReadMsgUDP.
     14 var udpOOBSize = func() int {
     15 	// We can't know whether we'll get an IPv4 control message or an
     16 	// IPv6 control message ahead of time. To get around this, we size
     17 	// the buffer equal to the largest of the two.
     18 
     19 	oob4 := ipv4.NewControlMessage(ipv4.FlagDst | ipv4.FlagInterface)
     20 	oob6 := ipv6.NewControlMessage(ipv6.FlagDst | ipv6.FlagInterface)
     21 
     22 	if len(oob4) > len(oob6) {
     23 		return len(oob4)
     24 	}
     25 
     26 	return len(oob6)
     27 }()
     28 
     29 // SessionUDP holds the remote address and the associated
     30 // out-of-band data.
     31 type SessionUDP struct {
     32 	raddr   *net.UDPAddr
     33 	context []byte
     34 }
     35 
     36 // RemoteAddr returns the remote network address.
     37 func (s *SessionUDP) RemoteAddr() net.Addr { return s.raddr }
     38 
     39 // ReadFromSessionUDP acts just like net.UDPConn.ReadFrom(), but returns a session object instead of a
     40 // net.UDPAddr.
     41 func ReadFromSessionUDP(conn *net.UDPConn, b []byte) (int, *SessionUDP, error) {
     42 	oob := make([]byte, udpOOBSize)
     43 	n, oobn, _, raddr, err := conn.ReadMsgUDP(b, oob)
     44 	if err != nil {
     45 		return n, nil, err
     46 	}
     47 	return n, &SessionUDP{raddr, oob[:oobn]}, err
     48 }
     49 
     50 // WriteToSessionUDP acts just like net.UDPConn.WriteTo(), but uses a *SessionUDP instead of a net.Addr.
     51 func WriteToSessionUDP(conn *net.UDPConn, b []byte, session *SessionUDP) (int, error) {
     52 	oob := correctSource(session.context)
     53 	n, _, err := conn.WriteMsgUDP(b, oob, session.raddr)
     54 	return n, err
     55 }
     56 
     57 func setUDPSocketOptions(conn *net.UDPConn) error {
     58 	// Try setting the flags for both families and ignore the errors unless they
     59 	// both error.
     60 	err6 := ipv6.NewPacketConn(conn).SetControlMessage(ipv6.FlagDst|ipv6.FlagInterface, true)
     61 	err4 := ipv4.NewPacketConn(conn).SetControlMessage(ipv4.FlagDst|ipv4.FlagInterface, true)
     62 	if err6 != nil && err4 != nil {
     63 		return err4
     64 	}
     65 	return nil
     66 }
     67 
     68 // parseDstFromOOB takes oob data and returns the destination IP.
     69 func parseDstFromOOB(oob []byte) net.IP {
     70 	// Start with IPv6 and then fallback to IPv4
     71 	// TODO(fastest963): Figure out a way to prefer one or the other. Looking at
     72 	// the lvl of the header for a 0 or 41 isn't cross-platform.
     73 	cm6 := new(ipv6.ControlMessage)
     74 	if cm6.Parse(oob) == nil && cm6.Dst != nil {
     75 		return cm6.Dst
     76 	}
     77 	cm4 := new(ipv4.ControlMessage)
     78 	if cm4.Parse(oob) == nil && cm4.Dst != nil {
     79 		return cm4.Dst
     80 	}
     81 	return nil
     82 }
     83 
     84 // correctSource takes oob data and returns new oob data with the Src equal to the Dst
     85 func correctSource(oob []byte) []byte {
     86 	dst := parseDstFromOOB(oob)
     87 	if dst == nil {
     88 		return nil
     89 	}
     90 	// If the dst is definitely an IPv6, then use ipv6's ControlMessage to
     91 	// respond otherwise use ipv4's because ipv6's marshal ignores ipv4
     92 	// addresses.
     93 	if dst.To4() == nil {
     94 		cm := new(ipv6.ControlMessage)
     95 		cm.Src = dst
     96 		oob = cm.Marshal()
     97 	} else {
     98 		cm := new(ipv4.ControlMessage)
     99 		cm.Src = dst
    100 		oob = cm.Marshal()
    101 	}
    102 	return oob
    103 }