82 lines
1.8 KiB
Go
82 lines
1.8 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"flag"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"log/slog"
|
||
|
|
"net"
|
||
|
|
"os"
|
||
|
|
"os/signal"
|
||
|
|
"sync"
|
||
|
|
"syscall"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func main() {
|
||
|
|
listenAddress := flag.String("listen", "127.0.0.1:10554", "local listen address")
|
||
|
|
upstreamAddress := flag.String("upstream", "", "upstream host:port")
|
||
|
|
flag.Parse()
|
||
|
|
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
|
||
|
|
if *upstreamAddress == "" {
|
||
|
|
logger.Error("upstream is required")
|
||
|
|
os.Exit(2)
|
||
|
|
}
|
||
|
|
if _, _, err := net.SplitHostPort(*upstreamAddress); err != nil {
|
||
|
|
logger.Error("upstream must be host:port")
|
||
|
|
os.Exit(2)
|
||
|
|
}
|
||
|
|
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
||
|
|
defer stop()
|
||
|
|
if err := serve(ctx, *listenAddress, *upstreamAddress); err != nil {
|
||
|
|
logger.Error("RTSP fault proxy stopped", "error", err)
|
||
|
|
os.Exit(1)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func serve(ctx context.Context, listenAddress, upstreamAddress string) error {
|
||
|
|
listener, err := net.Listen("tcp", listenAddress)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("listen: %w", err)
|
||
|
|
}
|
||
|
|
defer listener.Close()
|
||
|
|
go func() {
|
||
|
|
<-ctx.Done()
|
||
|
|
_ = listener.Close()
|
||
|
|
}()
|
||
|
|
var connections sync.WaitGroup
|
||
|
|
defer connections.Wait()
|
||
|
|
for {
|
||
|
|
client, acceptErr := listener.Accept()
|
||
|
|
if acceptErr != nil {
|
||
|
|
if ctx.Err() != nil {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
return fmt.Errorf("accept: %w", acceptErr)
|
||
|
|
}
|
||
|
|
connections.Add(1)
|
||
|
|
go func() {
|
||
|
|
defer connections.Done()
|
||
|
|
proxy(client, upstreamAddress)
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func proxy(client net.Conn, upstreamAddress string) {
|
||
|
|
defer client.Close()
|
||
|
|
upstream, err := net.DialTimeout("tcp", upstreamAddress, 5*time.Second)
|
||
|
|
if err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
defer upstream.Close()
|
||
|
|
done := make(chan struct{}, 2)
|
||
|
|
copyOneWay := func(destination, source net.Conn) {
|
||
|
|
_, _ = io.Copy(destination, source)
|
||
|
|
done <- struct{}{}
|
||
|
|
}
|
||
|
|
go copyOneWay(upstream, client)
|
||
|
|
go copyOneWay(client, upstream)
|
||
|
|
<-done
|
||
|
|
}
|