Optimize the memory footprint under SOCKS5 when USERNAME/PASSWORD authentication is used

[cty123]

Currently the Go bytes.Buffer struct is not used optimally, ref: https://github.com/cty123/tun2socks/blob/main/transport/socks5/socks5.go#L190

authMsg := &bytes.Buffer{}
authMsg.WriteByte(0x01 /* VER */)
authMsg.WriteByte(byte(len(user.Username)) /* ULEN */)
authMsg.WriteString(user.Username /* UNAME */)
authMsg.WriteByte(byte(len(user.Password)) /* PLEN */)
authMsg.WriteString(user.Password /* PASSWD */)

bytes.Buffer{} always returns an empty buffer with capacity = 0, and this capacity grows as data is written into the buffer. Under the hood, bytes.Buffer{} has a priavte grow(int) function to allocate additional space to store the data.

func (b *Buffer) Write(p []byte) (n int, err error) {
    b.lastRead = opInvalid
    m, ok := b.tryGrowByReslice(len(p))
    if !ok {
        m = b.grow(len(p))
    }
    return copy(b.buf[m:], p), nil
}

...

func (b *Buffer) grow(n int) int {
    m := b.Len()
    // If buffer is empty, reset to recover space.
    if m == 0 && b.off != 0 {
        b.Reset()
    }
    // Try to grow by means of a reslice.
    if i, ok := b.tryGrowByReslice(n); ok {
        return i
    }
    if b.buf == nil && n <= smallBufferSize {
        b.buf = make([]byte, n, smallBufferSize)
        return 0
    }
    c := cap(b.buf)
...

Therefore, during the operation to construct the AuthMsg payload, there are 2 things that are not optimal,

1. The buffer can grow multiple times as we are gradually putting more data into it, eg. username, password.
2. The capacity could be very large compared with the actual data we are storing.

What we can do instead is that, since the total amount of payload size can be calculated before the allocation, we can compute it and allocate a block of memory that fits exactly to the payload data, such that we are not wasting any memory and can avoid resizing the buffer during the handling.

For example,

package main

import (
    "bytes"
    "fmt"
)

func main() {
    username := "a very very very very very very very very very very very long username"
    password := "a very very very very very very very very very very very long password"

    fmt.Println("before")

    var authMsg *bytes.Buffer

    authMsg = &bytes.Buffer{}
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteByte(0x01 /* VER */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteByte(byte(len(username)) /* ULEN */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteString(username /* UNAME */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteByte(byte(len(password)) /* PLEN */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteString(password /* PASSWD */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    // cap: 0 len: 0
    // cap: 64 len: 1
    // cap: 64 len: 2
    // cap: 128 len: 72
    // cap: 128 len: 73
    // cap: 256 len: 143

    fmt.Println("after")

    authMsgLen := 1 + 1 + len(username) + 1 + len(password)
    authMsg = bytes.NewBuffer(make([]byte, 0, authMsgLen))
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteByte(0x01 /* VER */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteByte(byte(len(username)) /* ULEN */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteString(username /* UNAME */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteByte(byte(len(password)) /* PLEN */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    authMsg.WriteString(password /* PASSWD */)
    fmt.Println("cap:", authMsg.Cap(), "len:", authMsg.Len())

    // cap: 143 len: 0
    // cap: 143 len: 1
    // cap: 143 len: 2
    // cap: 143 len: 72
    // cap: 143 len: 73
    // cap: 143 len: 143
}

[xjasonlyu]

LGTM. Thanks for your contribution!