batch.go

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"runtime"
	"time"

	"github.com/klauspost/compress/zstd"
)

func (tun *tunnelCon) write(buf []byte) error {
	if tun == nil {
		return fmt.Errorf("nil tunnel")
	}

	tun.packetLock.Lock()
	defer tun.packetLock.Unlock()

	if tun.con == nil {
		return fmt.Errorf("tunnel is disconnected")
	}

	if batchingMicroseconds == 0 {
		tun.packets = buf
		tun.packetsLength = len(buf)
		return writeBatch(tun)
	} else {
		tun.packets = append(tun.packets, buf...)
		tun.packetsLength += len(buf)
		return nil
	}
}

func (tun *tunnelCon) batchWriter() {
	interval := time.Second
	if batchingMicroseconds > 0 {
		interval = time.Microsecond * time.Duration(batchingMicroseconds)
	}
	ticker := time.NewTicker(interval)

	if debugLog {
		defer doLog("batchWriter: exit")
	}

	for range ticker.C {
		if tun == nil || tun.con == nil {
			return
		}
		if time.Since(tun.lastUsed) > tunnelLife {
			doLog("Idle, disconnected")
			tun.delete(true)
			return
		}
		tun.packetLock.Lock()
		err := writeBatch(tun)
		tun.packetLock.Unlock()
		if err != nil {
			return
		}
	}

}

func writeBatch(tun *tunnelCon) error {
	if tun.packetsLength == 0 {
		return nil
	}
	if tun.con == nil {
		return fmt.Errorf("writeBatch: tunnel is disconnected")
	}

	var dataToWrite []byte
	if compressionLevel > 0 {
		var err error
		dataToWrite, err = compressFrame(tun)
		if err != nil {
			return err
		}
	} else {
		dataToWrite = tun.packets
	}

	var header []byte
	header = binary.AppendUvarint(header, uint64(len(dataToWrite)))
	frame := append(header, dataToWrite...)
	if err := writeFull(tun.con, frame); err != nil {
		return fmt.Errorf("writeBatch: %v", err)
	}

	tun.packets = nil
	tun.packetsLength = 0
	return nil
}

func compressFrame(tun *tunnelCon) ([]byte, error) {
	var buf bytes.Buffer
	zw, err := zstd.NewWriter(&buf, zstd.WithEncoderLevel(zstd.EncoderLevel(compressionLevel)), zstd.WithEncoderConcurrency(runtime.NumCPU()))
	if err != nil {
		return nil, fmt.Errorf("compressFrame: zstd writer error: %v", err)
	}
	if _, err := zw.Write(tun.packets); err != nil {
		_ = zw.Close()
		return nil, fmt.Errorf("compressFrame: zstd write error: %v", err)
	}
	if err := zw.Close(); err != nil {
		return nil, fmt.Errorf("compressFrame: zstd close error: %v", err)
	}
	return buf.Bytes(), nil
}

func decompressFrame(data []byte) ([]byte, error) {
	buf := bytes.NewReader(data)
	zr, err := zstd.NewReader(buf, zstd.WithDecoderConcurrency(runtime.NumCPU()))
	if err != nil {
		return nil, fmt.Errorf("decompressFrame: zstd reader error: %v", err)
	}
	defer zr.Close()

	var out bytes.Buffer
	if _, err := io.Copy(&out, zr); err != nil {
		return nil, fmt.Errorf("decompressFrame: zstd decompress copy error: %v", err)
	}
	return out.Bytes(), nil
}

func writeFull(con net.Conn, buf []byte) error {
	for len(buf) > 0 {
		n, err := con.Write(buf)
		if err != nil {
			return fmt.Errorf("write error: %v", err)
		}
		if n == 0 {
			return fmt.Errorf("short write: wrote 0 of %d bytes", len(buf))
		}
		buf = buf[n:]
	}
	return nil
}