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)
|
return co.Conn.Write(p)
|
||||||
}
|
}
|
||||||
|
|
||||||
l := make([]byte, 2)
|
msg := make([]byte, 2+len(p))
|
||||||
binary.BigEndian.PutUint16(l, uint16(len(p)))
|
binary.BigEndian.PutUint16(msg, uint16(len(p)))
|
||||||
|
copy(msg[2:], p)
|
||||||
n, err := (&net.Buffers{l, p}).WriteTo(co.Conn)
|
return co.Conn.Write(msg)
|
||||||
return int(n), err
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Return the appropriate timeout for a specific request
|
// 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) {
|
func TestTruncatedMsg(t *testing.T) {
|
||||||
m := new(Msg)
|
m := new(Msg)
|
||||||
m.SetQuestion("miek.nl.", TypeSRV)
|
m.SetQuestion("miek.nl.", TypeSRV)
|
||||||
|
|
|
@ -752,11 +752,10 @@ func (w *response) Write(m []byte) (int, error) {
|
||||||
return 0, &Error{err: "message too large"}
|
return 0, &Error{err: "message too large"}
|
||||||
}
|
}
|
||||||
|
|
||||||
l := make([]byte, 2)
|
msg := make([]byte, 2+len(m))
|
||||||
binary.BigEndian.PutUint16(l, uint16(len(m)))
|
binary.BigEndian.PutUint16(msg, uint16(len(m)))
|
||||||
|
copy(msg[2:], m)
|
||||||
n, err := (&net.Buffers{l, m}).WriteTo(w.tcp)
|
return w.tcp.Write(msg)
|
||||||
return int(n), err
|
|
||||||
default:
|
default:
|
||||||
panic("dns: internal error: udp and tcp both nil")
|
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 {
|
type ExampleFrameLengthWriter struct {
|
||||||
Writer
|
Writer
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue