tweaks for running very many copies

cmd/swarm/main.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log"
 	"math/rand"
+	"net/http"
 	"os"
 	"strconv"
 	"strings"
@@ -19,6 +20,54 @@ var s = rand.NewSource(time.Now().UnixNano()) // initialize global pseudo random
 var r = rand.New(s)
 var channel string
 var victim string
+var is_connected = false
+var greetChance = 100
+var suffixLength = 0
+
+func getNumericVar(key string, default_ int) int {
+	varS := os.Getenv(key)
+	if varS != "" {
+		var err error
+		default_, err = strconv.Atoi(varS)
+		if err != nil {
+			panic(fmt.Errorf("invalid %s: '%s'", key, varS))
+		}
+	}
+	return default_
+}
+
+func genNick() string {
+	return insults[r.Intn(len(insults))] + trailingDigits(suffixLength)
+}
+
+func trailingDigits(length int) string {
+	trailing := ""
+	for ; length > 0 ; length -- {
+		trailing = fmt.Sprintf("%s%d", trailing, r.Intn(10))
+	}
+	return trailing
+}
+
+func runStatusHandler() {
+	statusPort := getNumericVar("HEALTH_PORT", 8000)
+	if statusPort < 0 || statusPort > 65535 {
+		panic(fmt.Errorf("invalid HEALTH_PORT: %d", statusPort))
+	}
+	healthAddress := fmt.Sprintf("0.0.0.0:%d", statusPort)
+	log.Printf("serving health on %s", healthAddress)
+	http.HandleFunc("/health", statusHandler())
+	log.Fatal(http.ListenAndServe(healthAddress, nil))
+}
+
+func statusHandler() func(http.ResponseWriter, *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		if !is_connected {
+			http.Error(w, "Not connected", http.StatusInternalServerError)
+			return
+		}
+		w.Write([]byte("ok"))
+	}
+}
 
 func HandleRequest(ctx context.Context) (string, error) {
 	// So i was too lazy to get rid of the boilerplate
@@ -35,20 +84,28 @@ func HandleRequest(ctx context.Context) (string, error) {
 	if victim == "" {
 		log.Fatal("Set IRC_VICTIM!")
 	}
+	suffixLength = getNumericVar("SUFFIX_DIGITS", 0)
+	if suffixLength < 0 || suffixLength > 10 {
+		panic(fmt.Errorf("invalid SUFFIX_DIGITS: %d", suffixLength))
+	}
 
-	dialer := proxy.FromEnvironment()
+	log.Println("dialing...")
 
-	conn, err := dialer.Dial("tcp", server)
+	dialCtx, cancelDial := context.WithTimeout(context.Background(), time.Second*10)
+	defer cancelDial()
+
+	conn, err := proxy.Dial(dialCtx, "tcp", server)
 	if err != nil {
 		log.Fatal(err)
 	}
 	hostname, _ := os.Hostname()
 	config := irc.ClientConfig{
-		Nick:    insults[r.Intn(len(insults))],
+		Nick:    genNick(),
 		User:    hostname,
 		Name:    "dabbing on " + victim,
 		Handler: irc.HandlerFunc(handleIRCCon), // See below for callback
 	}
+	log.Println("starting client...")
 	client := irc.NewClient(conn, config)
 	err = client.Run()
 	return fmt.Sprintf("i died"), err
@@ -60,7 +117,7 @@ func handleIRCCon(c *irc.Client, m *irc.Message) {
 	// Only one reader, prob not necessary
 	// Handle connection established
 	if m.Command == "001" {
-
+		is_connected = true
 		// 001 is welcome event
 		c.Write("JOIN " + channel)
 
@@ -76,12 +133,12 @@ func handleIRCCon(c *irc.Client, m *irc.Message) {
 				},
 			})
 		} else if text[len(text)-1] == "name" {
-			nick := insults[r.Intn(len(insults))]
+			nick := genNick()
 			c.Write("NICK " + nick)
 		} else if text[len(text)-1] == "invite" {
 			c.Write("INVITE " + victim + " " + channel)
 			sent++
-		}/* else if strings.Contains(m.Trailing(), "greetings") {
+		} /* else if strings.Contains(m.Trailing(), "greetings") {
 			c.WriteMessage(&irc.Message{
 				Command: "PRIVMSG",
 				Params: []string{
@@ -92,18 +149,28 @@ func handleIRCCon(c *irc.Client, m *irc.Message) {
 		}*/
 	} else if m.Command == "JOIN" && c.FromChannel(m) {
 		if m.Prefix.Name != c.CurrentNick() {
-			c.WriteMessage(&irc.Message{
-				Command: "PRIVMSG",
-				Params: []string{
-					m.Params[0],
-					"greetings " + m.Prefix.Name,
-				},
-			})
+			if greetChance == 100 || r.Intn(100) < greetChance {
+				c.WriteMessage(&irc.Message{
+					Command: "PRIVMSG",
+					Params: []string{
+						m.Params[0],
+						"greetings " + m.Prefix.Name,
+					},
+				})
+			}
 		}
 	}
+
+	log.Printf("%+v\n", m)
 }
 
 func main() {
+	greetChance = getNumericVar("GREET_CHANCE", 100)
+	if greetChance > 100 || greetChance < 0 {
+		log.Fatalf("invalid GREET_CHANCE: '%s'", greetChance)
+	}
+
+	go runStatusHandler()
 	ctx, _ := context.WithCancel(context.Background())
 	ret, err := HandleRequest(ctx)
 	if err != nil {