From 9a4587957cd6989c71fe70d1ef6afe9baedb2d7d Mon Sep 17 00:00:00 2001
From: Akmal <72781956+Wikidepia@users.noreply.github.com>
Date: Mon, 27 Jan 2025 15:26:47 +0700
Subject: [PATCH] Add auth to remote scraper

---
 handlers/scraper/remote.go | 29 +++++++++++++++++++++++------
 main.go                    | 15 +++++++++++++++
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/handlers/scraper/remote.go b/handlers/scraper/remote.go
index 00c954e..8e704c5 100644
--- a/handlers/scraper/remote.go
+++ b/handlers/scraper/remote.go
@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"bytes"
 	"errors"
 	"log/slog"
 	"net"
@@ -20,25 +21,41 @@ type remoteResult struct {
 var sessCount atomic.Int32
 var inChan chan remoteResult
 
-func init() {
+func InitRemoteScraper(listenAddr *net.TCPAddr, authCode []byte) error {
+	if len(authCode) > 8 {
+		return errors.New("auth code max length is 8 bytes")
+	}
+
 	inChan = make(chan remoteResult)
 
-	ln, err := net.Listen("tcp", "0.0.0.0:4444")
+	ln, err := net.ListenTCP("tcp", listenAddr)
 	if err != nil {
-		return
+		return err
 	}
 	slog.Info("remote scraper is listening on", "address", ln.Addr())
 
-	go func() {
+	go func(ln *net.TCPListener, authCode []byte) {
 		for {
 			conn, err := ln.Accept()
 			if err != nil {
-				return
+				conn.Close()
+				continue
+			}
+
+			// deadline for read 5s
+			conn.SetReadDeadline(time.Now().Add(5 * time.Second))
+
+			authBytes := make([]byte, 8)
+			n, err := conn.Read(authBytes)
+			if err != nil || !bytes.Equal(authBytes[:n], authCode) {
+				conn.Close()
+				continue
 			}
 
 			go handleConnection(conn)
 		}
-	}()
+	}(ln, authCode)
+	return err
 }
 
 func handleConnection(conn net.Conn) {
diff --git a/main.go b/main.go
index 5ace274..ff1c7c0 100644
--- a/main.go
+++ b/main.go
@@ -7,6 +7,7 @@ import (
 	"instafix/utils"
 	"instafix/views"
 	"log/slog"
+	"net"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
@@ -28,6 +29,8 @@ func main() {
 	listenAddr := flag.String("listen", "0.0.0.0:3000", "Address to listen on")
 	gridCacheMaxFlag := flag.String("grid-cache-entries", "1024", "Maximum number of grid images to cache")
 	videoProxyAddr := flag.String("video-proxy-addr", "", "Video proxy address (https://github.com/Wikidepia/InstaFix-proxy)")
+	remoteAddr := flag.String("remote-scraper-addr", "0.0.0.0:3001", "Address of remote scraper")
+	authCode := flag.String("remote-scraper-auth-code", "", "Authentication code for remote scraper")
 	flag.Parse()
 
 	// Initialize video proxy
@@ -55,6 +58,18 @@ func main() {
 	scraper.InitDB()
 	defer scraper.DB.Close()
 
+	// Initialize remote scraper
+	if len(*authCode) != 0 {
+		remoteTCPAddr, err := net.ResolveTCPAddr("tcp", *remoteAddr)
+		if err != nil {
+			panic(err)
+		}
+		err = scraper.InitRemoteScraper(remoteTCPAddr, []byte(*authCode))
+		if err != nil {
+			panic(err)
+		}
+	}
+
 	// Evict cache every minute
 	go func() {
 		for {