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

2025-05-15 18:09:38 +10:00
parent b645f2b6bd
commit 169e0539f6
7 changed files with 472 additions and 374 deletions
--- a/testlib/client/client.go
+++ b/testlib/client/client.go
@@ -0,0 +1,184 @@
+package client
+
+/*
+   Copyright 2025 Suyono <suyono3484@gmail.com>
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"time"
+
+	"gitea.suyono.dev/suyono/netbounce/slicewriter"
+	"github.com/rs/zerolog/log"
+	"github.com/spf13/viper"
+)
+
+const (
+	SERVER      = "server"
+	PROTOCOL    = "protocol"
+	UDP         = "udp"
+	TCP         = "tcp"
+	NAME        = "name"
+	MESSAGE     = "message"
+	SLEEP       = "sleep"
+	READTIMEOUT = "read-timeout"
+)
+
+var gLimit *counter
+
+func SendMessages(ctx context.Context) {
+	sendMessages(ctx)
+}
+
+func InitLimit() {
+	gLimit = makeCounter(viper.GetInt("number"))
+}
+
+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])
+	}
+}
+
+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
+}