Skip to content

Conversation

cty123
Copy link
Contributor

@cty123 cty123 commented Nov 12, 2023

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
}

…hentication is used

* Optimize the usage of `bytes.Buffer` to avoid memory resizing.
* Rename magic number to constants for better reusability.
* Add tests for SOCKS5 request handling.
@xjasonlyu xjasonlyu self-requested a review November 12, 2023 16:25
Copy link
Owner

@xjasonlyu xjasonlyu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your PR! Could you address those changes before merging?

@cty123 cty123 requested a review from xjasonlyu November 12, 2023 21:47
@xjasonlyu
Copy link
Owner

LGTM. Thanks for your contribution!

@xjasonlyu xjasonlyu merged commit e86b3b7 into xjasonlyu:main Nov 12, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants