proxy/pkg/utils/utils.go

package utils

import (
	"context"
	"io"
	"log/slog"
	"net"
	"sync"

	"github.com/pkg/errors"
)

func ReadByte(reader io.Reader) (byte, error) {
	buffer, err := ReadBuffer(reader, 1)
	if err != nil {
		return 0, err
	}

	return buffer[0], nil
}

func ReadBuffer(reader io.Reader, size int) ([]byte, error) {
	buffer := make([]byte, size)
	_, err := io.ReadFull(reader, buffer)
	if err != nil {
		return nil, err
	}

	return buffer, nil
}

// Close 关闭对象，传入值绝对不能为 nil
func Close[T io.Closer](v T) {
	err := v.Close()
	if err != nil {
		slog.Warn("对象关闭失败", "err", err)
	}
}

func ConnChan(ctx context.Context, ls net.Listener) chan net.Conn {
	connCh := make(chan net.Conn)
	go func() {
		for {
			conn, err := ls.Accept()
			if err != nil {
				slog.Error("接受连接失败", err)
				// 临时错误重试连接
				var ne net.Error
				if errors.As(err, &ne) && ne.Temporary() {
					slog.Debug("临时错误重试")
					continue
				}
				return
			}
			// ctx 取消后退出
			select {
			case <-ctx.Done():
				Close(conn)
				return
			case connCh <- conn:
			}
		}
	}()
	return connCh
}

func WaitChan(ctx context.Context, wg *sync.WaitGroup) chan struct{} {
	ch := make(chan struct{})
	go func() {
		wg.Wait()
		select {
		case <-ctx.Done():
		case ch <- struct{}{}:
		}
	}()
	return ch
}