From f939294723f32ae62f629d4a4b836b9f763fd174 Mon Sep 17 00:00:00 2001
From: June <june@jsts.xyz>
Date: Sat, 16 Aug 2025 06:01:30 +0200
Subject: [PATCH] improve and streamline config options and introduce
 command-line flags

Introduce command-line flags for config options with the added benefit
of getting a help output, expose more config options to be directly
configurable, prefix all environment variables with "MRVC", adjust
config option naming to better reflect their purpose and add validation
to ensure some config options got explicitly set.
---
 main.go | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 72 insertions(+), 7 deletions(-)

diff --git a/main.go b/main.go
index e028382..1884cf8 100644
--- a/main.go
+++ b/main.go
@@ -2,11 +2,13 @@ package main
 
 import (
 	"context"
+	"flag"
 	"fmt"
 	"log"
 	"os"
 	"slices"
 	"sort"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -34,6 +36,12 @@ type VersionPath struct {
 	Version        string
 }
 
+var userIdFlag = flag.String("user-id", "", "The Matrix user ID to use. (EnvVar: MRVC_USER_ID)")
+var tokenFlag = flag.String("token", "", "The Matrix access token to use. (EnvVar: MRVC_TOKEN)")
+var roomFlag = flag.String("room", "", "The Matrix room to check. (EnvVar: MRVC_ROOM)")
+var printHomeserverMemberCountFlag = flag.Bool("print-homeserver-member-count", false, "Print the member count for each homeserver. (EnvVar: MRVC_PRINT_HOMESERVER_MEMBER_COUNT)")
+var homeserverVersionInfoTimeoutFlag = flag.Duration("homeserver-version-info-timeout", time.Second*5, "Timeout for getting the homeservers version information per homeserver. (EnvVar: MRVC_HOMESERVER_VERSION_INFO_TIMEOUT)")
+
 var unknownServerVersionInfo = fclient.Version{
 	Server: struct {
 		Name    string `json:"name"`
@@ -281,11 +289,68 @@ func compareVersionStrings(a, b string) int {
 }
 
 func main() {
-	userIdString := os.Getenv("MATRIX_USER_ID")
-	token := os.Getenv("MATRIX_TOKEN")
-	room := os.Getenv("MATRIX_ROOM")
-	verbose := false
-	timeoutSeconds := 5
+	flag.Parse()
+
+	// Configuration variables.
+	var userIdString, token, room string
+	var printHomeserverMemberCount bool
+	var homeserverVersionInfoTimeout time.Duration
+
+	visitedFlags := make(map[string]bool)
+	flag.Visit(func(visitedFlag *flag.Flag) {
+		visitedFlags[visitedFlag.Name] = true
+	})
+
+	// Assign flag and environment variable values to configuration variables.
+	// Flags take precedence over environment variables.
+	// This also ensures some configuration options got explicitly set.
+	if visitedFlags["user-id"] {
+		userIdString = *userIdFlag
+	} else if userIdEnvVar, ok := os.LookupEnv("MRVC_USER_ID"); ok {
+		userIdString = userIdEnvVar
+	} else {
+		log.Fatal("A Matrix user ID must be provided.")
+	}
+	if visitedFlags["token"] {
+		token = *tokenFlag
+	} else if tokenEnvVar, ok := os.LookupEnv("MRVC_TOKEN"); ok {
+		token = tokenEnvVar
+	} else {
+		log.Fatal("A Matrix access token must be provided.")
+	}
+	if visitedFlags["room"] {
+		room = *roomFlag
+	} else if roomEnvVar, ok := os.LookupEnv("MRVC_ROOM"); ok {
+		room = roomEnvVar
+	} else {
+		log.Fatal("A Matrix room must be provided.")
+	}
+	if visitedFlags["print-homeserver-member-count"] {
+		printHomeserverMemberCount = *printHomeserverMemberCountFlag
+	} else if printHomeserverMemberCountEnvVar, ok := os.LookupEnv("MRVC_PRINT_HOMESERVER_MEMBER_COUNT"); ok {
+		parsedPrintHomeserverMemberCountEnvVar, err := strconv.ParseBool(printHomeserverMemberCountEnvVar)
+		if err != nil {
+			log.Println("Error parsing MRVC_PRINT_HOMESERVER_MEMBER_COUNT:")
+			log.Fatal(err)
+		}
+		printHomeserverMemberCount = parsedPrintHomeserverMemberCountEnvVar
+	} else {
+		// The flag holds the default value.
+		printHomeserverMemberCount = *printHomeserverMemberCountFlag
+	}
+	if visitedFlags["homeserver-version-info-timeout"] {
+		homeserverVersionInfoTimeout = *homeserverVersionInfoTimeoutFlag
+	} else if homeserverVersionInfoTimeoutEnvVar, ok := os.LookupEnv("MRVC_HOMESERVER_VERSION_INFO_TIMEOUT"); ok {
+		parsedHomeserverVersionInfoTimeoutEnvVar, err := time.ParseDuration(homeserverVersionInfoTimeoutEnvVar)
+		if err != nil {
+			log.Println("Error parsing MRVC_HOMESERVER_VERSION_INFO_TIMEOUT:")
+			log.Fatal(err)
+		}
+		homeserverVersionInfoTimeout = parsedHomeserverVersionInfoTimeoutEnvVar
+	} else {
+		// The flag holds the default value.
+		homeserverVersionInfoTimeout = *homeserverVersionInfoTimeoutFlag
+	}
 
 	userId := id.UserID(userIdString)
 	_, homeserver, err := userId.ParseAndValidate()
@@ -302,7 +367,7 @@ func main() {
 	}
 	federationClient := fclient.NewClient(
 		fclient.WithWellKnownSRVLookups(true),
-		fclient.WithTimeout(time.Second*time.Duration(timeoutSeconds)),
+		fclient.WithTimeout(homeserverVersionInfoTimeout),
 	)
 
 	joinedMembers, err := client.JoinedMembers(context.Background(), id.RoomID(room))
@@ -393,7 +458,7 @@ func main() {
 
 				fmt.Printf("    %s -> %d\n", versionKey, membersByVersionPath[VersionPath{maxRoomVersionKey, serverKey, versionKey}])
 
-				if verbose {
+				if printHomeserverMemberCount {
 					homeserverKeys := make([]string, 0, len(versionValue))
 					for key := range versionValue {
 						homeserverKeys = append(homeserverKeys, key)