Use net.Buffers for writing TCP message (#934)
This commit is contained in:
parent
1a5555c783
commit
337216f9a7
21
client.go
21
client.go
|
@ -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
15
server.go
|
@ -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")
|
||||
|
|
Loading…
Reference in New Issue