Connect to a TCP socket, receive a stream from it, and write that data to disk

socket_saver.go

package main

import (
	"flag"
	"log"
	"fmt"
	"os"
	"os/signal"
	"io"
	"context"
	"net"
	"net/url"
)

func signalWaiter(ctx context.Context, cancel context.CancelFunc) {
	doneSignal := make(chan os.Signal, 1)
	signal.Notify(doneSignal, os.Interrupt)

	select {
	case <-doneSignal:
	case <-ctx.Done():
	}

	defer cancel()
}

func saveSocket(ctx context.Context, host string, outputFile string) {
	var bytesWritten int64 = 0
	f, err := os.Create(outputFile)
	if err != nil {
		log.Printf("Failed to create new file `%s`: %v\n", outputFile, err)
		return
	}
	defer func() {
		log.Printf("Wrote `%d` bytes to file: %s\n", bytesWritten, outputFile)
		f.Close()
	}()
	log.Printf("File `%s` created.\n", outputFile)

	log.Printf("Dialing TCP `%s`...\n", host)
	conn, err := net.Dial("tcp", host)
	if err != nil {
		log.Printf("Failed to connect to server: %v", err)
		return
	}
	defer conn.Close()
	log.Printf("Connected! Beginning read...")

	go func() {
		<-ctx.Done()
		conn.Close()
	}()

	bytesWritten, err = io.Copy(f, conn)
	if err != nil {
		log.Printf("Error during copy: %v", err)
	}
}

func main() {
	tcpUrl := flag.String("url", "", "TCP URL to read from.")
	outputFile := flag.String("o", "dump.ts", "Output file path")
	flag.Parse()

	if tcpUrl == nil || len(*tcpUrl) == 0 {
		log.Fatal("You must specify `--url`.")
	}

	u, err := url.Parse(*tcpUrl)
	if err != nil {
		log.Fatal(fmt.Errorf("Failed to parse URL `%s`: %w", err))
	}

	if u.Scheme != "tcp" {
		log.Fatal("Your URL is not a TCP url.")
	}

	ctx, cancel := context.WithCancel(context.Background())
	go signalWaiter(ctx, cancel)

	saveSocket(ctx, u.Host, *outputFile)
}