WIP: moved logic code in test tools to separate package for better code architecture

This commit is contained in:
2025-05-15 18:09:38 +10:00
parent b645f2b6bd
commit 169e0539f6
7 changed files with 472 additions and 374 deletions

View File

@@ -1,40 +0,0 @@
package slicewriter
import (
"sync"
"testing"
)
func TestDummy(t *testing.T) {
var (
p sync.Pool
a, b []byte
)
p = sync.Pool{
New: func() any {
return make([]byte, 4096)
},
}
a = p.Get().([]byte)
t.Logf("retrieve a from pool: len %d, cap %d", len(a), cap(a))
b = p.Get().([]byte)
t.Logf("retrieve b from pool: len %d, cap %d", len(b), cap(b))
a = a[:1024]
b = b[:1024]
t.Logf("resize a : len %d, cap %d", len(a), cap(a))
t.Logf("resize b : len %d, cap %d", len(b), cap(b))
p.Put(a[:cap(a)])
p.Put(b)
t.Log("after putting back")
a = p.Get().([]byte)
t.Logf("retrieve a from pool: len %d, cap %d", len(a), cap(a))
b = p.Get().([]byte)
t.Logf("retrieve b from pool: len %d, cap %d", len(b), cap(b))
}

View File

@@ -1,36 +0,0 @@
package slicewriter
import "io"
type SliceWriter struct {
off int
buf []byte
}
func (w *SliceWriter) Write(p []byte) (int, error) {
if len(p)+w.off <= len(w.buf) {
copy(w.buf[w.off:], p)
w.off += len(p)
return len(p), nil
}
space := len(w.buf) - w.off
copy(w.buf[:w.off], p[:space])
w.off += space
return space, io.ErrShortWrite
}
func (w *SliceWriter) Reset() {
w.off = 0
}
func (w *SliceWriter) Bytes() []byte {
return w.buf[:w.off]
}
func NewSliceWriter(b []byte) *SliceWriter {
return &SliceWriter{
off: 0,
buf: b,
}
}

View File

@@ -18,185 +18,64 @@ package main
import (
"context"
"fmt"
"net"
"os/signal"
"time"
"golang.org/x/sys/unix"
"gitea.suyono.dev/suyono/netbounce/cmd/slicewriter"
"gitea.suyono.dev/suyono/netbounce/testlib/client"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"github.com/spf13/pflag"
"github.com/spf13/viper"
)
var gLimit *counter
func main() {
ctx, cancel := signal.NotifyContext(context.Background(), unix.SIGINT, unix.SIGTERM)
defer cancel()
parseFlags()
log.Debug().Caller().Msgf("Sending messages to %s server: %s", viper.GetString(PROTOCOL), viper.GetString(SERVER))
sendMessages(ctx)
}
const (
SERVER = "server"
PROTOCOL = "protocol"
UDP = "udp"
TCP = "tcp"
NAME = "name"
MESSAGE = "message"
SLEEP = "sleep"
READTIMEOUT = "read-timeout"
)
func sendMessages(ctx context.Context) {
switch viper.GetString(PROTOCOL) {
case "udp":
sendUDP(ctx)
case "tcp":
sendTCP(ctx)
default:
log.Fatal().Caller().Str(PROTOCOL, viper.GetString(PROTOCOL)).Msg("Unknown protocol")
}
}
func sendTCP(ctx context.Context) {
var (
conn net.Conn
err error
buf, b []byte
n int
)
if conn, err = net.Dial(TCP, viper.GetString(SERVER)); err != nil {
log.Fatal().Caller().Err(err).Msg("Failed to connect to server")
}
defer func() {
_ = conn.Close()
}()
go func() {
<-ctx.Done()
_ = conn.Close()
}()
buf = make([]byte, 4096)
for gLimit.isContinue(ctx) {
sb := slicewriter.NewSliceWriter(buf)
if _, err = fmt.Fprintf(sb, "client %s | %v | %s", viper.GetString(NAME), time.Now(), viper.GetString(MESSAGE)); err != nil {
log.Fatal().Caller().Err(err).Msg("Failed to build client message")
}
b = sb.Bytes()
if _, err = conn.Write(b); err != nil {
log.Fatal().Caller().Err(err).Str(PROTOCOL, TCP).Str(SERVER, viper.GetString(SERVER)).Msg("Failed to send client message")
}
if err = conn.SetReadDeadline(time.Now().Add(viper.GetDuration(READTIMEOUT))); err != nil {
log.Fatal().Caller().Err(err).Str(PROTOCOL, TCP).Str(SERVER, viper.GetString(SERVER)).Msg("Failed to send client message")
}
if n, err = conn.Read(buf); err != nil {
log.Fatal().Caller().Err(err).Str(PROTOCOL, TCP).Str(SERVER, viper.GetString(SERVER)).Msg("read from the server")
}
log.Info().Caller().Str(PROTOCOL, TCP).Str(SERVER, viper.GetString(SERVER)).Msgf("%s", buf[:n])
time.Sleep(viper.GetDuration(SLEEP))
}
}
func sendUDP(ctx context.Context) {
var (
addr, laddr *net.UDPAddr
conn *net.UDPConn
rAddr net.Addr
err error
buf, b []byte
n int
)
if addr, err = net.ResolveUDPAddr(UDP, viper.GetString(SERVER)); err != nil {
log.Fatal().Caller().Err(err).Str(SERVER, viper.GetString(SERVER)).Msg("udp resolve server address")
}
if laddr, err = net.ResolveUDPAddr(UDP, ""); err != nil {
log.Fatal().Caller().Err(err).Str(SERVER, viper.GetString(SERVER)).Msg("udp resolve local/self address")
}
log.Info().Str(SERVER, viper.GetString(SERVER)).Msgf("bound address %v", laddr)
// In Go, binding address and port for UDP use ListenUDP. Confusing!!
if conn, err = net.ListenUDP(UDP, laddr); err != nil {
log.Fatal().Caller().Err(err).Str(SERVER, viper.GetString(SERVER)).Msg("fail to bind local/self address")
}
defer func() {
_ = conn.Close()
}()
go func() {
<-ctx.Done()
_ = conn.Close()
}()
buf = make([]byte, 4096)
for gLimit.isContinue(ctx) {
sb := slicewriter.NewSliceWriter(buf)
if _, err = fmt.Fprintf(sb, "client %s | %v | %s", viper.GetString(NAME), time.Now(), viper.GetString(MESSAGE)); err != nil {
log.Fatal().Caller().Err(err).Msg("Failed to build client message")
}
b = sb.Bytes()
if _, err = conn.WriteTo(b, addr); err != nil {
log.Fatal().Caller().Err(err).Str(PROTOCOL, UDP).Str(SERVER, viper.GetString(SERVER)).Msg("Failed to send client message")
}
if err = conn.SetReadDeadline(time.Now().Add(viper.GetDuration(READTIMEOUT))); err != nil {
log.Error().Caller().Err(err).Str(PROTOCOL, UDP).Str(SERVER, viper.GetString(SERVER)).Msg("set read timeout on the socket")
}
if n, rAddr, err = conn.ReadFrom(b); err != nil {
log.Fatal().Caller().Err(err).Str(PROTOCOL, UDP).Str(SERVER, viper.GetString(SERVER)).Msg("read from server")
}
log.Info().Caller().Str(PROTOCOL, UDP).Str(SERVER, rAddr.String()).Msgf("%s", buf[:n])
}
log.Debug().Caller().Msgf("Sending messages to %s server: %s", viper.GetString(client.PROTOCOL), viper.GetString(client.SERVER))
// sendMessages(ctx)
client.SendMessages(ctx)
}
func parseFlags() {
pflag.String(NAME, "", "client name")
pflag.String(client.NAME, "", "client name")
pflag.IntP("number", "n", 5, "number of messages to send; default 5; set to 0 for infinite")
pflag.Duration(READTIMEOUT, 5*time.Second, "read timeout; default 5 seconds")
pflag.DurationP(SLEEP, "s", 10*time.Millisecond, "sleep time between requests; default 10ms")
pflag.StringP(MESSAGE, "m", "message from client", "message to send")
pflag.Duration(client.READTIMEOUT, 5*time.Second, "read timeout; default 5 seconds")
pflag.DurationP(client.SLEEP, "s", 10*time.Millisecond, "sleep time between requests; default 10ms")
pflag.StringP(client.MESSAGE, "m", "message from client", "message to send")
pflag.Bool("debug", false, "run in debug mode")
pflag.String(TCP, "", "tcp server address")
pflag.String(UDP, "", "udp server address")
pflag.String(client.TCP, "", "tcp server address")
pflag.String(client.UDP, "", "udp server address")
pflag.Parse()
_ = viper.BindPFlags(pflag.CommandLine)
gLimit = makeCounter(viper.GetInt("number"))
// gLimit = makeCounter(viper.GetInt("number"))
client.InitLimit()
if !viper.IsSet(NAME) {
if !viper.IsSet(client.NAME) {
log.Fatal().Caller().Msg("server name is required")
}
if viper.IsSet(TCP) && viper.IsSet(UDP) {
if viper.IsSet(client.TCP) && viper.IsSet(client.UDP) {
log.Fatal().Caller().Msg("cannot use tcp and udp at once")
}
if !viper.IsSet(TCP) && !viper.IsSet(UDP) {
if !viper.IsSet(client.TCP) && !viper.IsSet(client.UDP) {
log.Fatal().Caller().Msg("--tcp or --udp is required, use one of them not both.")
}
if viper.IsSet(TCP) {
viper.Set(PROTOCOL, TCP)
viper.Set(SERVER, viper.GetString(TCP))
if viper.IsSet(client.TCP) {
viper.Set(client.PROTOCOL, client.TCP)
viper.Set(client.SERVER, viper.GetString(client.TCP))
}
if viper.IsSet(UDP) {
viper.Set(PROTOCOL, UDP)
viper.Set(SERVER, viper.GetString(UDP))
if viper.IsSet(client.UDP) {
viper.Set(client.PROTOCOL, client.UDP)
viper.Set(client.SERVER, viper.GetString(client.UDP))
}
zerolog.SetGlobalLevel(zerolog.InfoLevel)
@@ -204,30 +83,3 @@ func parseFlags() {
zerolog.SetGlobalLevel(zerolog.DebugLevel)
}
}
type counter struct {
limit, tick int
}
func makeCounter(limit int) *counter {
if limit <= 0 {
log.Fatal().Msg("number must be > 0")
}
return &counter{limit: limit, tick: -1}
}
func (c *counter) isContinue(ctx context.Context) bool {
select {
case <-ctx.Done():
return false
default:
}
if c.limit == 0 {
return true
}
c.tick++
return c.tick < c.limit
}

View File

@@ -18,13 +18,10 @@ package main
import (
"context"
"fmt"
"net"
"os/signal"
"sync"
"time"
"gitea.suyono.dev/suyono/netbounce/cmd/slicewriter"
"gitea.suyono.dev/suyono/netbounce/testlib/server"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"github.com/spf13/pflag"
@@ -41,7 +38,7 @@ func main() {
defer cancel()
wg := &sync.WaitGroup{}
openPorts(wg, ctx)
server.OpenPorts(wg, ctx)
wg.Wait()
@@ -71,167 +68,3 @@ func parseFlags() {
zerolog.SetGlobalLevel(zerolog.DebugLevel)
}
}
func openPorts(wg *sync.WaitGroup, ctx context.Context) {
tcpPorts := viper.GetStringSlice("tcp")
for _, tcpPort := range tcpPorts {
wg.Add(1)
go listen(ctx, wg, tcpPort)
}
udpPorts := viper.GetStringSlice("udp")
for _, udpPort := range udpPorts {
wg.Add(1)
go bindUDP(ctx, wg, udpPort)
}
}
func ClosePacket(ctx context.Context, conn net.PacketConn) {
<-ctx.Done()
_ = conn.Close()
}
func bindUDP(ctx context.Context, wg *sync.WaitGroup, address string) {
defer wg.Done()
var (
conn net.PacketConn
err error
buf, b []byte
n int
addr net.Addr
)
log.Debug().Caller().Msgf("binding socket for UDP on %v", address)
if conn, err = net.ListenPacket("udp", address); err != nil {
panic(fmt.Errorf("failed to bind udp address: %v", err))
}
go ClosePacket(ctx, conn)
buf = make([]byte, 4096)
udpLoop:
for {
if n, addr, err = conn.ReadFrom(buf); err != nil && n == 0 {
select {
case <-ctx.Done():
break udpLoop
default:
}
log.Error().Err(err).Msg("failed to read packet")
continue udpLoop
}
log.Info().Str("client", addr.String()).Msgf("received message: %s", buf[:n])
select {
case <-ctx.Done():
break udpLoop
default:
}
sb := slicewriter.NewSliceWriter(buf)
if _, err = fmt.Fprintf(sb, "server: %s | UDP | %v", viper.GetString("name"), time.Now()); err != nil {
log.Error().Err(err).Msg("build server message")
}
b = sb.Bytes()
if n, err = conn.WriteTo(b, addr); err != nil {
select {
case <-ctx.Done():
break udpLoop
default:
}
log.Error().Err(err).Str("client", addr.String()).Msg("failed to write packet")
continue udpLoop
}
if n != len(b) {
log.Debug().Str("client", addr.String()).Msg("incomplete packet sent")
}
log.Info().Str("client", addr.String()).Msg("packet received and replied")
}
}
func CloseListener(ctx context.Context, listener net.Listener) {
<-ctx.Done()
_ = listener.Close()
}
func listen(ctx context.Context, wg *sync.WaitGroup, address string) {
defer wg.Done()
var (
listener net.Listener
err error
conn net.Conn
)
log.Debug().Caller().Msgf("listeng for TCP on %v", address)
if listener, err = net.Listen("tcp", address); err != nil {
log.Error().Err(err).Str("address", address).Msg("failed to listen")
return
}
go CloseListener(ctx, listener)
tcpIncoming:
for {
if conn, err = listener.Accept(); err != nil {
select {
case <-ctx.Done():
break tcpIncoming
default:
}
log.Error().Err(err).Str("address", address).Msg("failed to accept connection")
continue tcpIncoming
}
wg.Add(1)
go handleTCP(ctx, wg, conn)
}
}
func CloseConnection(ctx context.Context, conn net.Conn) {
<-ctx.Done()
_ = conn.Close()
}
func handleTCP(ctx context.Context, wg *sync.WaitGroup, conn net.Conn) {
defer wg.Done()
defer func() {
_ = conn.Close()
}()
var (
buf, b []byte
err error
n int
)
buf = make([]byte, 4096)
addr := conn.RemoteAddr()
cctx, cancel := context.WithCancel(ctx)
go CloseConnection(cctx, conn)
defer cancel()
for {
if n, err = conn.Read(buf); err != nil {
log.Error().Err(err).Str("client", addr.String()).Msg("failed to read data from TCP connection")
return
}
log.Info().Str("client", addr.String()).Msgf("received message: %s", buf[:n])
sb := slicewriter.NewSliceWriter(buf)
if _, err = fmt.Fprintf(sb, "server: %s | TCP | %v", viper.GetString("name"), time.Now()); err != nil {
log.Error().Err(err).Msg("build server message")
}
b = sb.Bytes()
if _, err = conn.Write(b); err != nil {
log.Error().Err(err).Str("client", addr.String()).Msg("failed to write data to TCP connection")
return
}
}
}