feat: allow overriding mirrors via environment variables

README.md

@@ -109,6 +109,13 @@ sudo chmod u+r /etc/apt/auth.conf.d/90ubuntu-advantage
 The location of the credentials can be configured using the environment variable
 `CHISEL_AUTH_DIR`.
 
+For custom Ubuntu mirrors, the following environment variables can be used to
+override default archive URLs:
+
+- `CHISEL_UBUNTU_REPO_ARCHIVE_URL` (default: `http://archive.ubuntu.com/ubuntu/`)
+- `CHISEL_UBUNTU_REPO_OLD_RELEASES_URL` (default: `http://old-releases.ubuntu.com/ubuntu/`)
+- `CHISEL_UBUNTU_REPO_PORTS_URL` (default: `http://ports.ubuntu.com/ubuntu-ports/`)
+
 ## Reference
 
 ### Chisel releases

internal/archive/archive.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"os"
 	"slices"
 	"strings"
 	"time"
@@ -201,13 +202,20 @@ func archiveURL(pro, arch string, oldRelease bool) (string, *credentials, error)
 	}
 
 	if oldRelease {
-		return ubuntuOldReleasesURL, nil, nil
+		return envOrDefault("CHISEL_UBUNTU_REPO_OLD_RELEASES_URL", ubuntuOldReleasesURL), nil, nil
 	}
 
 	if arch == "amd64" || arch == "i386" {
-		return ubuntuURL, nil, nil
+		return envOrDefault("CHISEL_UBUNTU_REPO_ARCHIVE_URL", ubuntuURL), nil, nil
 	}
-	return ubuntuPortsURL, nil, nil
+	return envOrDefault("CHISEL_UBUNTU_REPO_PORTS_URL", ubuntuPortsURL), nil, nil
+}
+
+func envOrDefault(key, defaultValue string) string {
+	if value := os.Getenv(key); value != "" {
+		return value
+	}
+	return defaultValue
 }
 
 func openUbuntu(options *Options) (Archive, error) {

internal/archive/archive_test.go

@@ -294,6 +294,67 @@ func (s *httpSuite) TestFetchPortsPackage(c *C) {
 	c.Assert(read(pkg), Equals, "mypkg4 1.4 data")
 }
 
+func (s *httpSuite) TestArchiveURLFromEnv(c *C) {
+	restoreArchive := fakeEnv("CHISEL_UBUNTU_REPO_ARCHIVE_URL", "")
+	defer restoreArchive()
+	restoreOldReleases := fakeEnv("CHISEL_UBUNTU_REPO_OLD_RELEASES_URL", "")
+	defer restoreOldReleases()
+	restorePorts := fakeEnv("CHISEL_UBUNTU_REPO_PORTS_URL", "")
+	defer restorePorts()
+
+	tests := []struct {
+		summary    string
+		envName    string
+		base       string
+		arch       string
+		oldRelease bool
+	}{{
+		summary: "Ubuntu archive URL override",
+		envName: "CHISEL_UBUNTU_REPO_ARCHIVE_URL",
+		base:    "http://mirror.example/ubuntu/",
+		arch:    "amd64",
+	}, {
+		summary: "Ubuntu ports URL override",
+		envName: "CHISEL_UBUNTU_REPO_PORTS_URL",
+		base:    "http://mirror.example/ubuntu-ports/",
+		arch:    "arm64",
+	}, {
+		summary:    "Ubuntu old releases URL override",
+		envName:    "CHISEL_UBUNTU_REPO_OLD_RELEASES_URL",
+		base:       "http://mirror.example/ubuntu-old/",
+		arch:       "amd64",
+		oldRelease: true,
+	}}
+
+	for _, test := range tests {
+		c.Logf("Summary: %s", test.summary)
+
+		err := os.Setenv(test.envName, test.base)
+		c.Assert(err, IsNil)
+
+		s.base = test.base
+		s.responses = make(map[string][]byte)
+		s.prepareArchive("jammy", "22.04", test.arch, []string{"main"})
+
+		options := archive.Options{
+			Label:      "ubuntu",
+			Version:    "22.04",
+			Arch:       test.arch,
+			Suites:     []string{"jammy"},
+			Components: []string{"main"},
+			CacheDir:   c.MkDir(),
+			PubKeys:    []*packet.PublicKey{s.pubKey},
+			OldRelease: test.oldRelease,
+		}
+
+		_, err = archive.Open(&options)
+		c.Assert(err, IsNil)
+
+		err = os.Setenv(test.envName, "")
+		c.Assert(err, IsNil)
+	}
+}
+
 func (s *httpSuite) TestFetchSecurityPackage(c *C) {
 
 	for i, suite := range []string{"jammy", "jammy-updates", "jammy-security"} {