Skip to content

Instantly share code, notes, and snippets.

@tanopwan
Last active June 1, 2023 23:02
Show Gist options
  • Star 3 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save tanopwan/0897b8b621c6405d014b1dd0c0febf12 to your computer and use it in GitHub Desktop.
Save tanopwan/0897b8b621c6405d014b1dd0c0febf12 to your computer and use it in GitHub Desktop.
How to peek the data from the reader stream ie. from TCP connection using bufio
package main
import (
"bufio"
"log"
"net"
"net/http"
"os"
"strings"
"time"
)
var (
connId = uint64(0)
)
func main() {
laddr, err := net.ResolveTCPAddr("tcp", ":9999")
if err != nil {
log.Printf("Failed to resolve local address: %s", err)
os.Exit(1)
}
listener, err := net.ListenTCP("tcp", laddr)
if err != nil {
log.Printf("Failed to open local port to listen: %s", err)
os.Exit(1)
}
for {
conn, err := listener.AcceptTCP()
if err != nil {
log.Printf("Failed to accept connection '%s'", err)
continue
}
connId++
go func(c *net.TCPConn) {
defer c.Close()
reader := bufio.NewReader(c)
// note that conn's reader will not be available,
// we have to assign a new reader to conn also
bb, err := reader.Peek(14)
if err != nil {
log.Printf("failed to pre-read protocol '%s'", err)
}
// can do anything with this reader
// below example is th check the protocol and return immediately
if strings.Contains(string(bb), "HTTP") { // GET / HTTP/1.1
log.Printf("found HTTP connection, return 200 OK")
p := "HTTP/1.1 200 OK\n"
p += "Date: " + time.Now().Format(http.TimeFormat) + "\n"
p += "Connection: close\n"
p += "\n"
c.Write([]byte(p))
} else { // or do something else
data, _, err := reader.ReadLine()
if err != nil {
log.Printf("failed to read TCP data '%s'", err)
}
c.Write([]byte("Hello from TCP echo server:"))
// you will be able to read the data that already peeked
c.Write(data)
}
}(conn)
}
}
@tanopwan
Copy link
Author

tanopwan commented Feb 4, 2021

Noted that conn's reader will not be available,

We have to assign a new reader to conn object
One way is wrapping it with a struct that comply the interfaces

type ReadWriteCloser interface {
	Reader
	Writer
	Closer
}
type conBuff struct {
	conn   *net.TCPConn
	reader *bufio.Reader
}

func (c *conBuff) Read(p []byte) (n int, err error) {
	return c.reader.Read(p)
}

func (c *conBuff) Write(p []byte) (n int, err error) {
	return c.conn.Write(p)
}

func (c *conBuff) Close() error {
	return c.conn.Close()
}
newConn := &conBuff{
	conn:   conn,
	reader: reader,
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment