Send DNS query in one packet when using TCP/TLS (#1219)
* Send DNS query in one packet when using TCP/TLS * fix review comments * Removed net.Buffers * Added unit-tests for writing messages over TCP in one go
This commit is contained in:
parent
7d5e1ea350
commit
67bd57debd
|
@ -340,11 +340,10 @@ func (co *Conn) Write(p []byte) (int, error) {
|
|||
return co.Conn.Write(p)
|
||||
}
|
||||
|
||||
l := make([]byte, 2)
|
||||
binary.BigEndian.PutUint16(l, uint16(len(p)))
|
||||
|
||||
n, err := (&net.Buffers{l, p}).WriteTo(co.Conn)
|
||||
return int(n), err
|
||||
msg := make([]byte, 2+len(p))
|
||||
binary.BigEndian.PutUint16(msg, uint16(len(p)))
|
||||
copy(msg[2:], p)
|
||||
return co.Conn.Write(msg)
|
||||
}
|
||||
|
||||
// Return the appropriate timeout for a specific request
|
||||
|
|
|
@ -396,6 +396,24 @@ func TestClientConn(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestClientConnWriteSinglePacket(t *testing.T) {
|
||||
c := &countingConn{}
|
||||
conn := Conn{
|
||||
Conn: c,
|
||||
}
|
||||
m := new(Msg)
|
||||
m.SetQuestion("miek.nl.", TypeTXT)
|
||||
err := conn.WriteMsg(m)
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("failed to write: %v", err)
|
||||
}
|
||||
|
||||
if c.writes != 1 {
|
||||
t.Fatalf("incorrect number of Write calls")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTruncatedMsg(t *testing.T) {
|
||||
m := new(Msg)
|
||||
m.SetQuestion("miek.nl.", TypeSRV)
|
||||
|
|
|
@ -752,11 +752,10 @@ func (w *response) Write(m []byte) (int, error) {
|
|||
return 0, &Error{err: "message too large"}
|
||||
}
|
||||
|
||||
l := make([]byte, 2)
|
||||
binary.BigEndian.PutUint16(l, uint16(len(m)))
|
||||
|
||||
n, err := (&net.Buffers{l, m}).WriteTo(w.tcp)
|
||||
return int(n), err
|
||||
msg := make([]byte, 2+len(m))
|
||||
binary.BigEndian.PutUint16(msg, uint16(len(m)))
|
||||
copy(msg[2:], m)
|
||||
return w.tcp.Write(msg)
|
||||
default:
|
||||
panic("dns: internal error: udp and tcp both nil")
|
||||
}
|
||||
|
|
|
@ -1100,6 +1100,37 @@ func TestResponseDoubleClose(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
type countingConn struct {
|
||||
net.Conn
|
||||
writes int
|
||||
}
|
||||
|
||||
func (c *countingConn) Write(p []byte) (int, error) {
|
||||
c.writes++
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func TestResponseWriteSinglePacket(t *testing.T) {
|
||||
c := &countingConn{}
|
||||
rw := &response{
|
||||
tcp: c,
|
||||
}
|
||||
rw.writer = rw
|
||||
|
||||
m := new(Msg)
|
||||
m.SetQuestion("miek.nl.", TypeTXT)
|
||||
m.Response = true
|
||||
err := rw.WriteMsg(m)
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("failed to write: %v", err)
|
||||
}
|
||||
|
||||
if c.writes != 1 {
|
||||
t.Fatalf("incorrect number of Write calls")
|
||||
}
|
||||
}
|
||||
|
||||
type ExampleFrameLengthWriter struct {
|
||||
Writer
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue