community
/
dns
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Use net.Buffers for writing TCP message (#934)

This commit is contained in:
Tom Thorogood 2019-03-11 00:16:14 +10:30 committed by Miek Gieben
parent 1a5555c783
commit 337216f9a7
2 changed files with 13 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
client.go
View File

@ -3,7 +3,6 @@ package dns
// A client implementation.
import (
"bytes"
"context"
"crypto/tls"
"encoding/binary"
@ -353,23 +352,19 @@ func (co *Conn) WriteMsg(m *Msg) (err error) {
// Write implements the net.Conn Write method.
func (co *Conn) Write(p []byte) (n int, err error) {
switch t := co.Conn.(type) {
switch co.Conn.(type) {
case *net.TCPConn, *tls.Conn:
w := t.(io.Writer)
lp := len(p)
if lp < 2 {
return 0, io.ErrShortBuffer
}
if lp > MaxMsgSize {
if len(p) > MaxMsgSize {
return 0, &Error{err: "message too large"}
}
l := make([]byte, 2, lp+2)
binary.BigEndian.PutUint16(l, uint16(lp))
p = append(l, p...)
n, err := io.Copy(w, bytes.NewReader(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
}
return co.Conn.Write(p)
}

15
server.go
View File

@ -3,7 +3,6 @@
package dns
import (
"bytes"
"context"
"crypto/tls"
"encoding/binary"
@ -701,18 +700,14 @@ func (w *response) Write(m []byte) (int, error) {
case w.udp != nil:
return WriteToSessionUDP(w.udp, m, w.udpSession)
case w.tcp != nil:
lm := len(m)
if lm < 2 {
return 0, io.ErrShortBuffer
}
if lm > MaxMsgSize {
if len(m) > MaxMsgSize {
return 0, &Error{err: "message too large"}
}
l := make([]byte, 2, 2+lm)
binary.BigEndian.PutUint16(l, uint16(lm))
m = append(l, m...)
n, err := io.Copy(w.tcp, bytes.NewReader(m))
l := make([]byte, 2)
binary.BigEndian.PutUint16(l, uint16(len(m)))
n, err := (&net.Buffers{l, m}).WriteTo(w.tcp)
return int(n), err
default:
panic("dns: internal error: udp and tcp both nil")