added a distroless build, made sure milla works with gvisor's runsc as the runtime

4 files changed, 92 insertions, 66 deletions

diff --git a/Dockerfile_distroless b/Dockerfile_distroless
new file mode 100644
index 0000000..5b7401b
--- /dev/null
+++ b/Dockerfile_distroless
@@ -0,0 +1,10 @@
+FROM golang:1.21 as builder
+WORKDIR /milla
+COPY go.sum go.mod /milla/
+RUN go mod download
+COPY *.go /milla/
+RUN CGO_ENABLED=0 go build
+
+FROM gcr.io/distroless/static-debian12
+COPY --from=builder /milla/milla "/usr/bin/milla"
+ENTRYPOINT ["milla"]
diff --git a/README.md b/README.md
index 222468a..724123f 100644
--- a/README.md
+++ b/README.md
@@ -65,6 +65,8 @@ The formatter to use. This tells chroma how to generate the color in the output.
 - `terminal16m` for treucolor terminals
 - `html` for HTML output
 
+**_NOTE_**: please note that the terminal formatters will increase the size of the IRC event. Depending on the IRC server, this may or may not be a problem.
+
 #### provider
 
 Which LLM provider to use. The supported options are:
@@ -83,7 +85,7 @@ The system message to use for ollama.
 
 #### clientCertPath
 
-The path to the client certificate to use for SASL external authentication.
+The path to the client certificate to use for client cert authentication.
 
 #### serverPass
 
@@ -137,6 +139,8 @@ Whether to use TLS to connect to the IRC server. This option is provided to supp
 
 #### disableSTSFallback
 
+Disables the "fallback" to a non-TLS connection if the strict transport policy expires and the first attempt to reconnect back to the TLS version fails.
+
 #### allowFlood
 
 Disable [girc](https://github.com/lrstanley/girc)'s built-in flood protection.
@@ -183,9 +187,21 @@ Get the value of all config options.
 
 Set a config option on the fly. Use the same name as the config file but capitalized.
 
+## Proxy Support
+
+milla will read and use the `ALL_PROXY` environment variable.
+It is rather a non-standard way of using the env var but you define your socks5 proxies like so:
+
+```
+ALL_PROXY=127.0.0.1:9050
+```
+
+**_NOTE_**: the proxy is used for making calls to the LLMs, not for connecting to the IRC server.
+
 ## Deploy
 
 An example docker compose file is provided in the repo under `docker-compose.yaml`.
+milla can be used with [gvisor](https://gvisor.dev/)'s docker runtime, `runsc`.
 
 ```yaml
 services:
@@ -197,26 +213,31 @@ services:
       resources:
         limits:
           memory: 64M
-    user: ${UID}:${GID}
     logging:
       driver: "json-file"
       options:
         max-size: "100m"
     networks:
       - millanet
+    user: 1000:1000
     restart: unless-stopped
     command: ["--config", "/opt/milla/config.toml"]
     volumes:
-      - ./config.toml:/opt/milla/config.toml
-      - /etc/ssl/certs:/etc/ssl/certs:ro
+      - ./config-gpt.toml:/opt/milla/config.toml
+      - /etc/localtime:/etc/localtime:ro
+      - /etc/resolv.conf:/etc/resolv.conf:ro
     cap_drop:
       - ALL
-    environment:
-      - SERVER_DEPLOYMENT_TYPE=deployment
+    runtime: runsc
 networks:
   millanet:
+    driver: bridge
 ```
 
+The env vars `UID`and `GID`need to be defined or they can replaces by your host user's uid and gid.<br/>
+
+As a convinience, there is a a [distroless](https://github.com/GoogleContainerTools/distroless) dockerfile, `Dockerfile_distroless` also provided.
+
 ## Thanks
 
 - [girc](https://github.com/lrstanley/girc)
diff --git a/go.mod b/go.mod
index e9bf270..f5a5b23 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/google/generative-ai-go v0.11.2
 	github.com/lrstanley/girc v0.0.0-20240125042120-9add3166e52e
 	github.com/sashabaranov/go-openai v1.19.3
+	golang.org/x/net v0.24.0
 	google.golang.org/api v0.176.1
 )
 
@@ -35,7 +36,6 @@ require (
 	go.opentelemetry.io/otel/metric v1.24.0 // indirect
 	go.opentelemetry.io/otel/trace v1.24.0 // indirect
 	golang.org/x/crypto v0.22.0 // indirect
-	golang.org/x/net v0.24.0 // indirect
 	golang.org/x/oauth2 v0.19.0 // indirect
 	golang.org/x/sync v0.7.0 // indirect
 	golang.org/x/sys v0.19.0 // indirect
diff --git a/main.go b/main.go
index b470738..68120e5 100644
--- a/main.go
+++ b/main.go
@@ -10,9 +10,9 @@ import (
 	"fmt"
 	"log"
 	"net/http"
-	"net/url"
 	"os"
 	"reflect"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
@@ -22,6 +22,7 @@ import (
 	"github.com/google/generative-ai-go/genai"
 	"github.com/lrstanley/girc"
 	openai "github.com/sashabaranov/go-openai"
+	"golang.org/x/net/proxy"
 	"google.golang.org/api/option"
 )
 
@@ -77,7 +78,6 @@ func NewTomlConfig() *TomlConfig {
 		ChromaStyle:         "rose-pine-moon",
 		ChromaFormatter:     "noop",
 		Provider:            "ollama",
-		ClientCertPath:      "milla.pem",
 		Temp:                0.5,  //nolint:gomnd
 		RequestTimeout:      10,   //nolint:gomnd
 		MillaReconnectDelay: 30,   //nolint:gomnd
@@ -137,26 +137,26 @@ func returnGeminiResponse(resp *genai.GenerateContentResponse) string {
 	return result
 }
 
-// func extractLast256ColorEscapeCode(str string) (string, error) {
-// 	pattern256F := `\033\[38;5;(\d+)m`
-// 	// pattern256B := `\033\[48;5;(\d+)m`
-// 	// pattern16mF := `\033\[38;2;(\d+);(\d+);(\d+)m`
-// 	// pattern16mB := `\033\[48;2;(\d+);(\d+);(\d+)m`
+func extractLast256ColorEscapeCode(str string) (string, error) {
+	pattern256F := `\033\[38;5;(\d+)m`
+	// pattern256B := `\033\[48;5;(\d+)m`
+	// pattern16mF := `\033\[38;2;(\d+);(\d+);(\d+)m`
+	// pattern16mB := `\033\[48;2;(\d+);(\d+);(\d+)m`
 
-// 	r, err := regexp.Compile(pattern256F)
-// 	if err != nil {
-// 		return "", fmt.Errorf("failed to compile regular expression: %w", err)
-// 	}
+	r, err := regexp.Compile(pattern256F)
+	if err != nil {
+		return "", fmt.Errorf("failed to compile regular expression: %w", err)
+	}
 
-// 	matches := r.FindAllStringSubmatch(str, -1)
-// 	if len(matches) == 0 {
-// 		return "", nil
-// 	}
+	matches := r.FindAllStringSubmatch(str, -1)
+	if len(matches) == 0 {
+		return "", nil
+	}
 
-// 	lastMatch := matches[len(matches)-1]
+	lastMatch := matches[len(matches)-1]
 
-// 	return lastMatch[1], nil
-// }
+	return lastMatch[1], nil
+}
 
 func chunker(inputString string, chromaFormatter string) []string {
 	chunks := strings.Split(inputString, "\n")
@@ -169,17 +169,16 @@ func chunker(inputString string, chromaFormatter string) []string {
 	case "terminal16":
 		fallthrough
 	case "terminal256":
-		// for count, chunk := range chunks {
-		// 	lastColorCode, err := extractLast256ColorEscapeCode(chunk)
-		// 	if err != nil {
-		// 		continue
-		// 	}
+		for count, chunk := range chunks {
+			lastColorCode, err := extractLast256ColorEscapeCode(chunk)
+			if err != nil {
+				continue
+			}
 
-		// 	if count <= len(chunks)-2 {
-		// 		chunks[count+1] = fmt.Sprintf("\033[38;5;%sm", lastColorCode) + chunks[count+1]
-		// 	}
-		// }
-		fallthrough
+			if count <= len(chunks)-2 {
+				chunks[count+1] = fmt.Sprintf("\033[38;5;%sm", lastColorCode) + chunks[count+1]
+			}
+		}
 	case "terminal16m":
 		fallthrough
 	default:
@@ -384,20 +383,14 @@ func ollamaHandler(
 
 		request.Header.Set("Content-Type", "application/json")
 
-		httpClient := http.Client{}
-		allProxy := os.Getenv("ALL_PROXY")
-		if allProxy != "" {
-			proxyURL, err := url.Parse(allProxy)
-			if err != nil {
-				client.Cmd.ReplyTo(event, fmt.Sprintf("error: %s", err.Error()))
+		var httpClient http.Client
 
-				return
-			}
-			transport := &http.Transport{
-				Proxy: http.ProxyURL(proxyURL),
-			}
+		dialer := proxy.FromEnvironment()
 
-			httpClient.Transport = transport
+		httpClient = http.Client{
+			Transport: &http.Transport{
+				Dial: dialer.Dial,
+			},
 		}
 
 		response, err := httpClient.Do(request)
@@ -564,24 +557,19 @@ func chatGPTHandler(
 		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(appConfig.RequestTimeout)*time.Second)
 		defer cancel()
 
-		allProxy := os.Getenv("ALL_PROXY")
-		config := openai.DefaultConfig(appConfig.Apikey)
-		if allProxy != "" {
-			proxyURL, err := url.Parse(allProxy)
-			if err != nil {
-				client.Cmd.ReplyTo(event, fmt.Sprintf("error: %s", err.Error()))
+		var httpClient http.Client
 
-				return
-			}
-			transport := &http.Transport{
-				Proxy: http.ProxyURL(proxyURL),
-			}
+		dialer := proxy.FromEnvironment()
 
-			config.HTTPClient = &http.Client{
-				Transport: transport,
-			}
+		httpClient = http.Client{
+			Transport: &http.Transport{
+				Dial: dialer.Dial,
+			},
 		}
 
+		config := openai.DefaultConfig(appConfig.Apikey)
+		config.HTTPClient = &httpClient
+
 		gptClient := openai.NewClientWithConfig(config)
 
 		*gptMemory = append(*gptMemory, openai.ChatCompletionMessage{
@@ -645,7 +633,7 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client) {
 		DisableSTSFallback: appConfig.DisableSTSFallback,
 		GlobalFormat:       true,
 		TLSConfig: &tls.Config{
-			InsecureSkipVerify: appConfig.SkipTLSVerify,
+			InsecureSkipVerify: appConfig.SkipTLSVerify, // #nosec G402
 			ServerName:         appConfig.IrcServer,
 		},
 	})
@@ -676,9 +664,16 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client) {
 		}
 	}
 
-	// if appConfig.EnableSasl && appConfig.ClientCertPath != "" {
-	// 	// TODO  - add client cert support
-	// }
+	if appConfig.EnableSasl && appConfig.ClientCertPath != "" {
+		cert, err := tls.LoadX509KeyPair(appConfig.ClientCertPath, appConfig.ClientCertPath)
+		if err != nil {
+			log.Println("invalid client certificate.")
+
+			return
+		}
+
+		irc.Config.TLSConfig.Certificates = []tls.Certificate{cert}
+	}
 
 	irc.Handlers.AddBg(girc.CONNECTED, func(c *girc.Client, e girc.Event) {
 		for _, channel := range appConfig.IrcChannels {
@@ -700,7 +695,7 @@ func runIRC(appConfig TomlConfig, ircChan chan *girc.Client) {
 	for {
 		if err := irc.Connect(); err != nil {
 			log.Println(err)
-			log.Println("reconnecting in" + strconv.Itoa(appConfig.MillaReconnectDelay))
+			log.Println("reconnecting in " + strconv.Itoa(appConfig.MillaReconnectDelay))
 			time.Sleep(time.Duration(appConfig.MillaReconnectDelay) * time.Second)
 		} else {
 			return