github.com/phuslu/fastdns@v0.8.3-0.20240310041952-69506fc67dd1/writer.go (about)

     1  package fastdns
     2  
     3  import (
     4  	"net"
     5  	"net/netip"
     6  )
     7  
     8  // A ResponseWriter interface is used by an DNS handler to construct an DNS response.
     9  type ResponseWriter interface {
    10  	// LocalAddr returns the netip.AddrPort of the server
    11  	LocalAddr() netip.AddrPort
    12  
    13  	// RemoteAddr returns the netip.AddrPort of the client that sent the current request.
    14  	RemoteAddr() netip.AddrPort
    15  
    16  	// Write writes a raw buffer back to the client.
    17  	Write([]byte) (int, error)
    18  }
    19  
    20  // MemResponseWriter is an implementation of ResponseWriter that supports write response to memory.
    21  type MemResponseWriter struct {
    22  	Data  []byte
    23  	Raddr netip.AddrPort
    24  	Laddr netip.AddrPort
    25  }
    26  
    27  // RemoteAddr returns the netip.AddrPort of the client that sent the current request.
    28  func (rw *MemResponseWriter) RemoteAddr() netip.AddrPort {
    29  	return rw.Raddr
    30  }
    31  
    32  // LocalAddr returns the netip.AddrPort of the server
    33  func (rw *MemResponseWriter) LocalAddr() netip.AddrPort {
    34  	return rw.Laddr
    35  }
    36  
    37  // Write writes a raw buffer back to the memory buffer.
    38  func (rw *MemResponseWriter) Write(p []byte) (n int, err error) {
    39  	rw.Data = append(rw.Data, p...)
    40  	n = len(p)
    41  	return
    42  }
    43  
    44  type udpResponseWriter struct {
    45  	Conn     *net.UDPConn
    46  	AddrPort netip.AddrPort
    47  }
    48  
    49  func (rw *udpResponseWriter) RemoteAddr() netip.AddrPort {
    50  	return rw.AddrPort
    51  }
    52  
    53  func (rw *udpResponseWriter) LocalAddr() netip.AddrPort {
    54  	return rw.Conn.LocalAddr().(*net.UDPAddr).AddrPort()
    55  }
    56  
    57  func (rw *udpResponseWriter) Write(p []byte) (n int, err error) {
    58  	n, _, err = rw.Conn.WriteMsgUDPAddrPort(p, nil, rw.AddrPort)
    59  	return
    60  }