WIP: moved logic code in test tools to separate package for better code architecture
This commit is contained in:
parent
b645f2b6bd
commit
169e0539f6
@ -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,
|
||||
}
|
||||
}
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -1,5 +1,21 @@
|
||||
package slicewriter
|
||||
|
||||
/*
|
||||
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 (
|
||||
"sync"
|
||||
"testing"
|
||||
52
slicewriter/slicewriter.go
Normal file
52
slicewriter/slicewriter.go
Normal file
@ -0,0 +1,52 @@
|
||||
package slicewriter
|
||||
|
||||
/*
|
||||
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 "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,
|
||||
}
|
||||
}
|
||||
184
testlib/client/client.go
Normal file
184
testlib/client/client.go
Normal file
@ -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
|
||||
}
|
||||
197
testlib/server/server.go
Normal file
197
testlib/server/server.go
Normal file
@ -0,0 +1,197 @@
|
||||
package server
|
||||
|
||||
/*
|
||||
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"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitea.suyono.dev/suyono/netbounce/slicewriter"
|
||||
"github.com/rs/zerolog/log"
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
func OpenPorts(wg *sync.WaitGroup, ctx context.Context) {
|
||||
openPorts(wg, ctx)
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user