Peek into Conn without reading in go

2019-04-11 16:26发布

问题:

I have a server net.Conn, and I'd like to peek into it before reading out bytes, to check if it's a plain text protocol the client is trying to use, or SSL/TLS.

Checking http://golang.org/pkg/net/, it seems the Conn interface does not have anything like that. I know I could use an iobuf.Reader, but I'd like to get a TLS Conn via tls.Conn(conn, config) if it turns out the client is using SSL/TLS, and iobuf.Reader would read from the original Conn, thus the handshake in tls.Conn would fail.

So is there any way to peek into a Conn in Go (something like MSG_PEEK in C/C++ sockets)? Or, to create a tls.Conn after I had read out the first few bytes from the underlying Conn?

回答1:

You're very close to a solution - the only thing you got wrong was reading from the Conn itself first. You are right that bufio.Reader's Peek method is the way to go. The trick is to make the buffered reader first and call Peek on the buffered reader rather than reading from the original Conn. Here's a bufferedConn type that will do what you need:

type bufferedConn struct {
    r        *bufio.Reader
    net.Conn // So that most methods are embedded
}

func newBufferedConn(c net.Conn) bufferedConn {
    return bufferedConn{bufio.NewReader(c), c}
}

func newBufferedConnSize(c net.Conn, n int) bufferedConn {
    return bufferedConn{bufio.NewReaderSize(c, n), c}
}

func (b bufferedConn) Peek(n int) ([]byte, error) {
    return b.r.Peek(n)
}

func (b bufferedConn) Read(p []byte) (int, error) {
    return b.r.Read(p)
}

What this does is allow you to access all of the normal net.Conn methods (by embedding the net.Conn - you could also write wrapper functions, but this is a lot easier and cleaner), and additionally provide access to the bufferedReader's Peek and Read methods (it's important that Read be called on the bufferedReader, not directly on the net.Conn because Peek stores data in a buffer, so subsequent calls to Read need to be able to first read any data out of this buffer before falling back to the underlying net.Conn).

The newBufferedConnSize function is probably unnecessary given that the current default buffer size is 4096 bytes, but technically if you're going to rely on being able to call Peek with a given size and not have it return an error (specifically ErrBufferFull), you should set it explicitly to a size that's at least as big as the number of bytes you intend to peek.

Check it out on the Go Playground.