Merge pull request #5732 from marxarelli/feature/http-accept

http: Support additional request headers

Author: tonistiigi

client/client_test.go

@@ -106,6 +106,7 @@ var allTests = []func(t *testing.T, sb integration.Sandbox){
 	testBuildHTTPSource,
 	testBuildHTTPSourceEtagScope,
 	testBuildHTTPSourceAuthHeaderSecret,
+	testBuildHTTPSourceHeader,
 	testBuildPushAndValidate,
 	testBuildExportWithUncompressed,
 	testBuildExportScratch,
@@ -3051,6 +3052,45 @@ func testBuildHTTPSourceAuthHeaderSecret(t *testing.T, sb integration.Sandbox) {
 	require.Equal(t, "Bearer foo", allReqs[0].Header.Get("Authorization"))
 }
 
+func testBuildHTTPSourceHeader(t *testing.T, sb integration.Sandbox) {
+	c, err := New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	modTime := time.Now().Add(-24 * time.Hour) // avoid falso positive with current time
+
+	resp := httpserver.Response{
+		Etag:         identity.NewID(),
+		Content:      []byte("content1"),
+		LastModified: &modTime,
+	}
+
+	server := httpserver.NewTestServer(map[string]httpserver.Response{
+		"/foo": resp,
+	})
+	defer server.Close()
+
+	st := llb.HTTP(
+		server.URL+"/foo",
+		llb.Header(llb.HTTPHeader{
+			Accept:    "application/vnd.foo",
+			UserAgent: "fooagent",
+		}),
+	)
+
+	def, err := st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.NoError(t, err)
+
+	allReqs := server.Stats("/foo").Requests
+	require.Equal(t, 1, len(allReqs))
+	require.Equal(t, http.MethodGet, allReqs[0].Method)
+	require.Equal(t, "application/vnd.foo", allReqs[0].Header.Get("accept"))
+	require.Equal(t, "fooagent", allReqs[0].Header.Get("user-agent"))
+}
+
 func testResolveAndHosts(t *testing.T, sb integration.Sandbox) {
 	requiresLinux(t)
 	c, err := New(sb.Context(), sb.Address())

client/llb/source.go

@@ -615,6 +615,10 @@ func HTTP(url string, opts ...HTTPOption) State {
 		attrs[pb.AttrHTTPAuthHeaderSecret] = hi.AuthHeaderSecret
 		addCap(&hi.Constraints, pb.CapSourceHTTPAuth)
 	}
+	if hi.Header != nil {
+		hi.Header.setAttrs(attrs)
+		addCap(&hi.Constraints, pb.CapSourceHTTPHeader)
+	}
 
 	addCap(&hi.Constraints, pb.CapSourceHTTP)
 	source := NewSource(url, attrs, hi.Constraints)
@@ -629,6 +633,7 @@ type HTTPInfo struct {
 	UID              int
 	GID              int
 	AuthHeaderSecret string
+	Header           *HTTPHeader
 }
 
 type HTTPOption interface {
@@ -666,6 +671,33 @@ func Chown(uid, gid int) HTTPOption {
 	})
 }
 
+// Header returns an [HTTPOption] that ensures additional request headers will
+// be sent when retrieving the HTTP source.
+func Header(header HTTPHeader) HTTPOption {
+	return httpOptionFunc(func(hi *HTTPInfo) {
+		hi.Header = &header
+	})
+}
+
+type HTTPHeader struct {
+	Accept    string
+	UserAgent string
+}
+
+func (hh *HTTPHeader) setAttrs(attrs map[string]string) {
+	if hh.Accept != "" {
+		attrs[hh.attr("accept")] = hh.Accept
+	}
+
+	if hh.UserAgent != "" {
+		attrs[hh.attr("user-agent")] = hh.UserAgent
+	}
+}
+
+func (hh *HTTPHeader) attr(name string) string {
+	return pb.AttrHTTPHeaderPrefix + name
+}
+
 func platformSpecificSource(id string) bool {
 	return strings.HasPrefix(id, "docker-image://") || strings.HasPrefix(id, "oci-layout://")
 }

solver/pb/attr.go

@@ -21,6 +21,7 @@ const AttrHTTPPerm = "http.perm"
 const AttrHTTPUID = "http.uid"
 const AttrHTTPGID = "http.gid"
 const AttrHTTPAuthHeaderSecret = "http.authheadersecret"
+const AttrHTTPHeaderPrefix = "http.header."
 
 const AttrImageResolveMode = "image.resolvemode"
 const AttrImageResolveModeDefault = "default"

solver/pb/caps.go

@@ -36,6 +36,7 @@ const (
 	CapSourceHTTPPerm     apicaps.CapID = "source.http.perm"
 	// NOTE the historical typo
 	CapSourceHTTPUIDGID apicaps.CapID = "soruce.http.uidgid"
+	CapSourceHTTPHeader apicaps.CapID = "source.http.header"
 
 	CapSourceOCILayout apicaps.CapID = "source.ocilayout"
 
@@ -238,13 +239,19 @@ func init() {
 	})
 
 	Caps.Init(apicaps.Cap{
-		ID:      CapSourceOCILayout,
+		ID:      CapSourceHTTPUIDGID,
 		Enabled: true,
 		Status:  apicaps.CapStatusExperimental,
 	})
 
 	Caps.Init(apicaps.Cap{
-		ID:      CapSourceHTTPUIDGID,
+		ID:      CapSourceHTTPHeader,
+		Enabled: true,
+		Status:  apicaps.CapStatusExperimental,
+	})
+
+	Caps.Init(apicaps.Cap{
+		ID:      CapSourceOCILayout,
 		Enabled: true,
 		Status:  apicaps.CapStatusExperimental,
 	})

source/http/identifier.go

@@ -26,6 +26,12 @@ type HTTPIdentifier struct {
 	UID              int
 	GID              int
 	AuthHeaderSecret string
+	Header           []HeaderField
+}
+
+type HeaderField struct {
+	Name  string
+	Value string
 }
 
 var _ source.Identifier = (*HTTPIdentifier)(nil)

source/http/source.go

@@ -2,6 +2,7 @@ package http
 
 import (
 	"bytes"
+	"cmp"
 	"context"
 	"crypto/sha256"
 	"encoding/json"
@@ -13,6 +14,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"slices"
 	"strconv"
 	"strings"
 	"time"
@@ -33,6 +35,13 @@ import (
 	"github.com/pkg/errors"
 )
 
+// supportedUserHeaders defines supported user-defined header fields. Fields
+// not included here will be silently dropped.
+var supportedUserDefinedHeaders = map[string]bool{
+	http.CanonicalHeaderKey("accept"):     true,
+	http.CanonicalHeaderKey("user-agent"): true,
+}
+
 type Opt struct {
 	CacheAccessor cache.Accessor
 	Transport     http.RoundTripper
@@ -95,9 +104,22 @@ func (hs *httpSource) Identifier(scheme, ref string, attrs map[string]string, pl
 			id.GID = int(i)
 		case pb.AttrHTTPAuthHeaderSecret:
 			id.AuthHeaderSecret = v
+		default:
+			if name, found := strings.CutPrefix(k, pb.AttrHTTPHeaderPrefix); found {
+				name = http.CanonicalHeaderKey(name)
+				if supportedUserDefinedHeaders[name] {
+					id.Header = append(id.Header, HeaderField{Name: name, Value: v})
+				}
+			}
 		}
 	}
 
+	// Sort header fields to ensure consistent hashing (see urlHash() and
+	// formatCacheKey())
+	slices.SortFunc(id.Header, func(a, b HeaderField) int {
+		return cmp.Compare(a.Name, b.Name)
+	})
+
 	return id, nil
 }
 
@@ -133,6 +155,7 @@ func (hs *httpSourceHandler) urlHash() (digest.Digest, error) {
 		Filename         []byte
 		Perm, UID, GID   int
 		AuthHeaderSecret string `json:",omitempty"`
+		Header           []HeaderField
 	}{
 		Filename: bytes.Join([][]byte{
 			[]byte(hs.src.URL),
@@ -142,6 +165,7 @@ func (hs *httpSourceHandler) urlHash() (digest.Digest, error) {
 		UID:              hs.src.UID,
 		GID:              hs.src.GID,
 		AuthHeaderSecret: hs.src.AuthHeaderSecret,
+		Header:           hs.src.Header,
 	})
 	if err != nil {
 		return "", err
@@ -154,8 +178,9 @@ func (hs *httpSourceHandler) formatCacheKey(filename string, dgst digest.Digest,
 		Filename         string
 		Perm, UID, GID   int
 		Checksum         digest.Digest
-		LastModTime      string `json:",omitempty"`
-		AuthHeaderSecret string `json:",omitempty"`
+		LastModTime      string        `json:",omitempty"`
+		AuthHeaderSecret string        `json:",omitempty"`
+		Header           []HeaderField `json:",omitempty"`
 	}{
 		Filename:         filename,
 		Perm:             hs.src.Perm,
@@ -164,6 +189,7 @@ func (hs *httpSourceHandler) formatCacheKey(filename string, dgst digest.Digest,
 		Checksum:         dgst,
 		LastModTime:      lastModTime,
 		AuthHeaderSecret: hs.src.AuthHeaderSecret,
+		Header:           hs.src.Header,
 	})
 	if err != nil {
 		return dgst
@@ -219,7 +245,7 @@ func (hs *httpSourceHandler) CacheKey(ctx context.Context, g session.Group, inde
 			for t := range m {
 				etags = append(etags, t)
 			}
-			req.Header.Add("If-None-Match", strings.Join(etags, ", "))
+			req.Header.Set("If-None-Match", strings.Join(etags, ", "))
 
 			if len(etags) == 1 {
 				onlyETag = etags[0]
@@ -236,7 +262,7 @@ func (hs *httpSourceHandler) CacheKey(ctx context.Context, g session.Group, inde
 		req.Method = "HEAD"
 		// we need to add accept-encoding header manually because stdlib only adds it to GET requests
 		// some servers will return different etags if Accept-Encoding header is different
-		req.Header.Add("Accept-Encoding", "gzip")
+		req.Header.Set("Accept-Encoding", "gzip")
 		resp, err := client.Do(req)
 		if err == nil {
 			if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusNotModified {
@@ -482,6 +508,9 @@ func (hs *httpSourceHandler) newHTTPRequest(ctx context.Context, g session.Group
 	}
 
 	req.Header.Set("User-Agent", version.UserAgent())
+	for _, field := range hs.src.Header {
+		req.Header.Set(field.Name, field.Value)
+	}
 
 	if hs.src.AuthHeaderSecret != "" {
 		err := hs.sm.Any(ctx, g, func(ctx context.Context, _ string, caller session.Caller) error {