suyono/netbounce
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Suyono 2025-05-15 18:09:38 +10:00
parent b645f2b6bd
commit 169e0539f6
7 changed files with 472 additions and 374 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
cmd/slicewriter/slicewriter.go
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,
}
}

190
cmd/test-client/main.go
View File

@ -18,185 +18,64 @@ package main
import ( import (
"context" "context"
"fmt"
"net"
"os/signal" "os/signal"
"time" "time"
"golang.org/x/sys/unix" "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"
"github.com/rs/zerolog/log" "github.com/rs/zerolog/log"
"github.com/spf13/pflag" "github.com/spf13/pflag"
"github.com/spf13/viper" "github.com/spf13/viper"
) )
var gLimit *counter
func main() { func main() {
ctx, cancel := signal.NotifyContext(context.Background(), unix.SIGINT, unix.SIGTERM) ctx, cancel := signal.NotifyContext(context.Background(), unix.SIGINT, unix.SIGTERM)
defer cancel() defer cancel()
parseFlags() parseFlags()
log.Debug().Caller().Msgf("Sending messages to %s server: %s", viper.GetString(PROTOCOL), viper.GetString(SERVER)) log.Debug().Caller().Msgf("Sending messages to %s server: %s", viper.GetString(client.PROTOCOL), viper.GetString(client.SERVER))
sendMessages(ctx) // sendMessages(ctx)
} client.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])
}
} }
func parseFlags() { 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.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.Duration(client.READTIMEOUT, 5*time.Second, "read timeout; default 5 seconds")
pflag.DurationP(SLEEP, "s", 10*time.Millisecond, "sleep time between requests; default 10ms") pflag.DurationP(client.SLEEP, "s", 10*time.Millisecond, "sleep time between requests; default 10ms")
pflag.StringP(MESSAGE, "m", "message from client", "message to send") pflag.StringP(client.MESSAGE, "m", "message from client", "message to send")
pflag.Bool("debug", false, "run in debug mode") pflag.Bool("debug", false, "run in debug mode")
pflag.String(TCP, "", "tcp server address") pflag.String(client.TCP, "", "tcp server address")
pflag.String(UDP, "", "udp server address") pflag.String(client.UDP, "", "udp server address")
pflag.Parse() pflag.Parse()
_ = viper.BindPFlags(pflag.CommandLine) _ = 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") 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") 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.") log.Fatal().Caller().Msg("--tcp or --udp is required, use one of them not both.")
} }
if viper.IsSet(TCP) { if viper.IsSet(client.TCP) {
viper.Set(PROTOCOL, TCP) viper.Set(client.PROTOCOL, client.TCP)
viper.Set(SERVER, viper.GetString(TCP)) viper.Set(client.SERVER, viper.GetString(client.TCP))
} }
if viper.IsSet(UDP) { if viper.IsSet(client.UDP) {
viper.Set(PROTOCOL, UDP) viper.Set(client.PROTOCOL, client.UDP)
viper.Set(SERVER, viper.GetString(UDP)) viper.Set(client.SERVER, viper.GetString(client.UDP))
} }
zerolog.SetGlobalLevel(zerolog.InfoLevel) zerolog.SetGlobalLevel(zerolog.InfoLevel)
@ -204,30 +83,3 @@ func parseFlags() {
zerolog.SetGlobalLevel(zerolog.DebugLevel) 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
}

171
cmd/test-server/main.go
View File

@ -18,13 +18,10 @@ package main
import ( import (
"context" "context"
"fmt"
"net"
"os/signal" "os/signal"
"sync" "sync"
"time"
"gitea.suyono.dev/suyono/netbounce/cmd/slicewriter" "gitea.suyono.dev/suyono/netbounce/testlib/server"
"github.com/rs/zerolog" "github.com/rs/zerolog"
"github.com/rs/zerolog/log" "github.com/rs/zerolog/log"
"github.com/spf13/pflag" "github.com/spf13/pflag"
@ -41,7 +38,7 @@ func main() {
defer cancel() defer cancel()
wg := &sync.WaitGroup{} wg := &sync.WaitGroup{}
openPorts(wg, ctx) server.OpenPorts(wg, ctx)
wg.Wait() wg.Wait()
@ -71,167 +68,3 @@ func parseFlags() {
zerolog.SetGlobalLevel(zerolog.DebugLevel) 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
}
}
}

16
cmd/slicewriter/dummy_test.go → slicewriter/dummy_test.go
View File

@ -1,5 +1,21 @@
package slicewriter 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 ( import (
"sync" "sync"
"testing" "testing"

52
slicewriter/slicewriter.go Normal file
View 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
View 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
View 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
}
}
}