tbin/tbin.go

package main

import (
	"crypto/sha256"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
)

var (
	connHost    = flag.String("host", "0.0.0.0", "Host to listen on")
	connPort    = flag.String("port", "4242", "Port to listen on")
	httpPort    = flag.String("http-port", "8080", "Port to listen on")
	httpURL     = flag.String("http-url", "http://localhost:8080/", "Base url to send back to the client")
	dir         = flag.String("dir", "./data", "Directory to store snippets in")
	buff        = flag.Int("buff", 8, "Buffer size in MB")
	chunkSize   = flag.Int("chunk", 8192, "Chunk size in bytes")
	connTimeout = flag.Int("conn-time", 10, "Connection timeout in seconds")
	readTimeout = flag.Int("read-time", 2, "Read timeout in seconds")
	ttl         = flag.Int("ttl", 48, "Time-to-live in hours for the snippets")
	verbose     = flag.Bool("verbose", false, "Verbose output")
	plaintext   = flag.Bool("plaintext", false, "Only allow plaintext snippets")
)

func main() {
	flag.Parse()

	// Validate and normalize configuration
	*buff *= 1024 * 1024
	if *buff <= 0 {
		log.Fatal("Buffer size must be positive")
	}

	if *chunkSize <= 0 {
		log.Fatal("Chunk size must be positive")
	}

	if *connTimeout <= 0 {
		log.Fatal("Connection timeout must be positive")
	}

	if *readTimeout <= 0 {
		log.Fatal("Read timeout must be positive")
	}

	if *ttl <= 0 {
		log.Fatal("TTL must be positive")
	}

	if !strings.HasSuffix(*httpURL, "/") {
		*httpURL += "/"
	}

	// Validate and create data directory
	if *dir == "" {
		log.Fatal("Data directory cannot be empty")
	}

	// Use absolute path for data directory to prevent path traversal
	absDir, err := filepath.Abs(*dir)
	if err != nil {
		log.Fatalf("Error getting absolute path for data directory: %v", err)
	}
	*dir = absDir

	if err := os.MkdirAll(*dir, 0700); err != nil {
		log.Fatalf("Error creating data directory %s: %v", *dir, err)
	}

	log.Printf("Starting tbin with data directory: %s", *dir)
	log.Printf("Buffer size: %d MB", *buff/1024/1024)
	log.Printf("TTL: %d hours", *ttl)

	go listenTCP()
	listenHTTP()
}

func listenTCP() {
	l, err := net.Listen("tcp", *connHost+":"+*connPort)
	if err != nil {
		log.Fatalf("Error listening: %v", err)
	}
	defer l.Close()
	log.Printf("Listening on %s:%s for TCP connections", *connHost, *connPort)
	for {
		conn, err := l.Accept()
		if err != nil {
			log.Printf("Error accepting: %v", err)
			continue
		}
		go handleTCP(conn)
	}
}

func listenHTTP() {
	// Set up security headers middleware
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Security headers
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("X-Frame-Options", "DENY")
		w.Header().Set("X-XSS-Protection", "1; mode=block")
		w.Header().Set("Referrer-Policy", "no-referrer")
		w.Header().Set("Content-Security-Policy", "default-src 'self'; script-src 'none'; style-src 'self' 'unsafe-inline'; img-src 'self'; font-src 'none'; connect-src 'none'; frame-src 'none'; object-src 'none'")

		handleHTTP(w, r)
	})

	log.Printf("Listening on %s:%s for HTTP connections", *connHost, *httpPort)
	err := http.ListenAndServe(*connHost+":"+*httpPort, nil)
	if err != nil {
		log.Fatalf("Error listening: %v", err)
	}
}

func handleHTTP(w http.ResponseWriter, r *http.Request) {
	if *verbose {
		log.Printf("HTTP %s %s %s", r.RemoteAddr, r.Method, r.URL.Path)
	}

	// Security: Prevent path traversal attacks
	id := strings.TrimPrefix(r.URL.Path, "/")
	if strings.Contains(id, "..") || strings.Contains(id, "/") {
		http.Error(w, "Invalid path", http.StatusBadRequest)
		return
	}

	switch r.Method {
	case "POST":
		switch id {
		case "paste":
			data := r.FormValue("paste")
			if data == "" {
				http.Error(w, "Empty data", http.StatusBadRequest)
				return
			}
			if len(data) > *buff {
				http.Error(w, "Data too large", http.StatusRequestEntityTooLarge)
				return
			}
			id = handleData([]byte(data))
			if id == "" {
				http.Error(w, "Error generating file", http.StatusInternalServerError)
				return
			}
			http.Redirect(w, r, "/"+id, http.StatusSeeOther)
		case "remove":
			id = r.FormValue("remove")
			if id != "" {
				go handleRemove(id)
			}
			http.Redirect(w, r, "/", http.StatusSeeOther)
		}
	case "GET":
		switch {
		case id == "":
			id = "index.html"
			w.Header().Set("Content-Type", "text/html")
		case strings.HasPrefix(id, "favicon"):
			id = "favicon.png"
			w.Header().Set("Content-Type", "image/png")
		case id == "example":
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.Write([]byte("Hello there\n"))
			return
		case id == "purge":
			// Security: Only allow purge from localhost or add authentication
			if r.RemoteAddr != "[::1]:" && !strings.HasPrefix(r.RemoteAddr, "127.0.0.1:") {
				log.Printf("Unauthorized purge attempt from %s", r.RemoteAddr)
				http.Error(w, "Unauthorized", http.StatusUnauthorized)
				return
			}
			log.Printf("Purge requested from %s", r.RemoteAddr)
			go handlePurge()
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.Write([]byte("OK\n"))
			return
		default:
			w.Header().Set("Content-Type", "text/plain; charset=utf-8")
			w.Header().Set("X-Robots-Tag", "noindex,nofollow")
			id = *dir + "/" + id + ".txt"
		}
		f, err := os.Open(id)
		if err != nil {
			http.Error(w, "Not found", http.StatusNotFound)
			return
		}
		defer f.Close()
		io.CopyN(w, f, int64(*buff))
	}
}

func handleTCP(conn net.Conn) {
	defer conn.Close()

	// Rate limiting: limit concurrent connections per IP
	remoteAddr := conn.RemoteAddr().String()
	log.Printf("TCP %s connection", remoteAddr)

	buf := make([]byte, *chunkSize)
	var data []byte
	var dataLen int

	conn.SetReadDeadline(time.Now().Add(time.Duration(*connTimeout) * time.Second))

	for {
		n, err := conn.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("Error reading from %s: %v", remoteAddr, err)
			}
			break
		}
		data = append(data, buf[:n]...)
		dataLen += n

		if dataLen >= *buff {
			log.Printf("Buffer limit exceeded from %s: %d bytes", remoteAddr, dataLen)
			conn.Write([]byte("Warning: " + strconv.Itoa(*buff/1024/1024) + "MB limit exceeded\n"))
			break
		}

		conn.SetReadDeadline(time.Now().Add(time.Duration(*readTimeout) * time.Second))
	}

	if dataLen == 0 {
		log.Printf("No data received from %s", remoteAddr)
		conn.Write([]byte("Error: no data received\n"))
		return
	}

	if *plaintext {
		err := checkData(data)
		if err != nil {
			log.Printf("Invalid data from %s: %v", remoteAddr, err)
			conn.Write([]byte("Error: " + err.Error() + "\n"))
			return
		}
	}

	id := handleData(data)
	if id == "" {
		log.Printf("Error generating file for %s", remoteAddr)
		conn.Write([]byte("Error generating file\n"))
		return
	}

	log.Printf("TCP %s %d %s.txt", remoteAddr, dataLen, id)
	conn.Write([]byte(*httpURL + id + "\n"))
	time.Sleep(100 * time.Millisecond)
}

func checkData(data []byte) error {
	if !utf8.Valid(data) {
		return errors.New("only plain text data is allowed")
	}

	// Additional validation: check for control characters (except newline)
	for _, b := range data {
		if b < 32 && b != '\n' && b != '\t' && b != '\r' {
			return errors.New("control characters are not allowed")
		}
	}

	return nil
}

func handleData(data []byte) string {
	var id string
	for l := 6; l <= 16; l++ {
		id = generateId(data, l)
		if _, err := os.Stat(*dir + "/" + id + ".txt"); os.IsNotExist(err) {
			break
		}
		id = ""
	}
	if id == "" {
		log.Println("Error no available id found")
		return ""
	}

	// Validate ID format
	if !isAlphanumeric(id) {
		log.Printf("Error invalid ID format: %s", id)
		return ""
	}

	filename := *dir + "/" + id + ".txt"
	f, err := os.Create(filename)
	if err != nil {
		log.Printf("Error creating file %s: %v", filename, err)
		return ""
	}
	defer f.Close()

	// Ensure data ends with newline
	if len(data) > 0 && data[len(data)-1] != '\n' {
		data = append(data, '\n')
	}

	_, err = f.Write(data)
	if err != nil {
		log.Printf("Error writing to file %s: %v", filename, err)
		// Try to clean up the empty file
		os.Remove(filename)
		return ""
	}

	return id
}

func generateId(data []byte, length int) string {
	hash := sha256.Sum256(data)
	id := fmt.Sprintf("%x", hash)
	if len(id) > length {
		id = id[:length]
	}
	return id
}

func isAlphanumeric(id string) bool {
	for _, c := range id {
		if !((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9')) {
			return false
		}
	}
	return true
}

func handleRemove(id string) {
	// Extract ID from URL or direct input
	id = strings.TrimPrefix(strings.Split(id, "/")[len(strings.Split(id, "/"))-1], "/")

	// Validate ID format
	if id == "" || len(id) > 16 || !isAlphanumeric(id) {
		log.Printf("Error invalid id to remove: %s", id)
		return
	}

	filename := *dir + "/" + id + ".txt"

	// Check if file exists before attempting to remove
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		log.Printf("File not found for removal: %s", filename)
		return
	}

	err := os.Remove(filename)
	if err != nil {
		log.Printf("Error removing file %s: %v", filename, err)
	} else {
		log.Printf("File removed: %s", filename)
	}
}

func handlePurge() {
	files, err := os.ReadDir(*dir)
	if err != nil {
		log.Printf("Error reading data directory: %v", err)
		return
	}

	var filesRemoved int
	var filesSkipped int
	var errorsEncountered int

	for _, file := range files {
		if file.IsDir() {
			filesSkipped++
			continue
		}

		// Only process .txt files
		if !strings.HasSuffix(file.Name(), ".txt") {
			filesSkipped++
			continue
		}

		info, err := file.Info()
		if err != nil {
			log.Printf("Error getting file info for %s: %v", file.Name(), err)
			errorsEncountered++
			continue
		}

		if time.Since(info.ModTime()) > time.Duration(*ttl)*time.Hour {
			filename := *dir + "/" + file.Name()
			err = os.Remove(filename)
			if err != nil {
				log.Printf("Error removing file %s: %v", filename, err)
				errorsEncountered++
			} else {
				filesRemoved++
				log.Printf("File removed: %s", filename)
			}
		}
	}

	log.Printf("Purge done: %d files removed, %d files skipped, %d errors", filesRemoved, filesSkipped, errorsEncountered)
}