tcp-stream-proxy/main.go

87 lines
1.6 KiB
Go
Raw Normal View History

2023-01-28 19:30:56 +01:00
package main
import (
"io"
"net"
"os"
"sync"
"github.com/integrii/flaggy"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
)
func main() {
log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
var (
listenAddr = ":8000"
forwardAddr string
)
flaggy.String(&listenAddr, "l", "listen", "Address to listen on")
flaggy.String(&forwardAddr, "f", "forward", "Address to forward to")
flaggy.Parse()
if forwardAddr == "" {
log.Fatal().Msg("Forward address is required")
}
ln, err := net.Listen("tcp", listenAddr)
if err != nil {
log.Panic().Err(err).Msg("Failed to listen")
}
defer ln.Close()
log.Info().Str("Listen Address", listenAddr).Msg("Listening")
for {
conn, err := ln.Accept()
if err != nil {
log.Err(err).Msg("Failed to accept connection")
continue
}
log.Info().Str("Incoming Address", conn.RemoteAddr().String()).Msg("Connection accepted")
go handleConnection(conn, forwardAddr)
}
}
func handleConnection(conn net.Conn, forwardAddr string) {
defer conn.Close()
fwrd, err := net.Dial("tcp", forwardAddr)
if err != nil {
log.Err(err).Msg("Failed to connect to forward address")
return
}
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer fwrd.Close()
defer wg.Done()
_, err := io.Copy(fwrd, conn)
if err != nil {
log.Err(err).Msg("Failed to copy from connection to forward address")
}
}()
go func() {
defer conn.Close()
defer wg.Done()
_, err := io.Copy(conn, fwrd)
if err != nil {
log.Err(err).Msg("Failed to copy from forward address to connection")
}
}()
wg.Wait()
log.Info().Str("Incoming Address", conn.RemoteAddr().String()).Msg("Connection closed")
}