feat: add authinfo injector (#2149)

Signed-off-by: Jim Ma <[email protected]>

client/dfget/dfget.go

@@ -227,7 +227,7 @@ func newDownRequest(cfg *config.DfgetConfig, hdr map[string]string) *dfdaemonv1.
 	} else {
 		rg = cfg.Range
 	}
-	return &dfdaemonv1.DownRequest{
+	request := &dfdaemonv1.DownRequest{
 		Url:               cfg.URL,
 		Output:            cfg.Output,
 		Timeout:           uint64(cfg.Timeout),
@@ -246,6 +246,19 @@ func newDownRequest(cfg *config.DfgetConfig, hdr map[string]string) *dfdaemonv1.
 		Gid:                int64(os.Getgid()),
 		KeepOriginalOffset: cfg.KeepOriginalOffset,
 	}
+
+	_url, err := url.Parse(cfg.URL)
+	if err == nil {
+		inj, ok := source.ShouldInjectAuthInfo(_url.Scheme)
+		if ok {
+			err = inj.Inject(_url, request.UrlMeta)
+			if err != nil {
+				logger.Errorf("inject auth info error: %s", err)
+			}
+		}
+	}
+
+	return request
 }
 
 func newProgressBar(max int64) *progressbar.ProgressBar {

pkg/source/clients/orasprotocol/oras_source_client.go

@@ -21,12 +21,14 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"net/http"
+	"net/url"
 	"os"
 	"regexp"
 	"strings"
 
+	commonv1 "d7y.io/api/pkg/apis/common/v1"
+
 	logger "d7y.io/dragonfly/v2/internal/dflog"
 	"d7y.io/dragonfly/v2/pkg/source"
 )
@@ -35,7 +37,9 @@ const (
 	ociAcceptHeader  = "application/vnd.oci.image.manifest.v1+json"
 	jsonAcceptHeader = "application/json"
 	scheme           = "oras"
-	authFilePath     = "/.singularity/docker-config.json"
+	configFilePath   = "/.singularity/docker-config.json"
+
+	authHeader = "X-Dragonfly-Oras-Authorization"
 )
 
 var _ source.ResourceClient = (*orasSourceClient)(nil)
@@ -49,7 +53,9 @@ type Manifest struct {
 }
 
 func init() {
-	source.RegisterBuilder(scheme, source.NewPlainResourceClientBuilder(Builder))
+	source.RegisterBuilder(scheme,
+		source.NewPlainResourceClientBuilder(Builder),
+		source.WithAuthInfoInjector(source.NewPlainAuthInfoInjector(AuthInfoInjector)))
 }
 
 func Builder(optionYaml []byte) (source.ResourceClient, source.RequestAdapter, []source.Hook, error) {
@@ -64,6 +70,16 @@ func Builder(optionYaml []byte) (source.ResourceClient, source.RequestAdapter, [
 	return client, client.adaptor, nil, nil
 }
 
+func AuthInfoInjector(_url *url.URL, urlMeta *commonv1.UrlMeta) error {
+	auth, err := fetchAuthInfo(_url.Host, false)
+	if err != nil {
+		return err
+	}
+
+	urlMeta.Header[authHeader] = "Basic " + auth
+	return nil
+}
+
 func (client *orasSourceClient) adaptor(request *source.Request) *source.Request {
 	clonedRequest := request.Clone(request.Context())
 	return clonedRequest
@@ -108,30 +124,53 @@ func (client *orasSourceClient) Download(request *source.Request) (*source.Respo
 	return imageFetchResponse, nil
 }
 
+func fetchAuthInfo(host string, skipCheckExist bool) (string, error) {
+	configFile := os.Getenv("HOME") + configFilePath
+	if !skipCheckExist && !fileExists(configFile) {
+		return "", nil
+	}
+
+	var auth string
+	databytes, err := os.ReadFile(configFile)
+	if err != nil {
+		return "", err
+	}
+
+	var jsonData map[string]interface{}
+	if err = json.Unmarshal(databytes, &jsonData); err != nil {
+		return "", err
+	}
+
+	for _, v := range jsonData {
+		for registry, v1 := range (v).(map[string]interface{}) {
+			for _, credentials := range (v1).(map[string]interface{}) {
+				if registry == host {
+					auth = credentials.(string)
+				}
+			}
+		}
+	}
+	return auth, nil
+}
+
 func (client *orasSourceClient) fetchToken(request *source.Request, path string) (string, error) {
 	var response *http.Response
 	var err error
 	tokenFetchURL := fmt.Sprintf("https://%s/service/token/?scope=repository:%s:pull&service=harbor-registry", request.URL.Host, path)
-	if fileExists(os.Getenv("HOME") + authFilePath) {
-		var auth string
-		databytes, err := os.ReadFile(os.Getenv("HOME") + authFilePath)
+	if authHeaderVal := request.Header.Get(authHeader); authHeaderVal != "" {
+		// remove the internal auth header
+		request.Header.Del(authHeader)
+		response, err = client.doRequest(request, jsonAcceptHeader, authHeaderVal, tokenFetchURL)
 		if err != nil {
 			return "", err
 		}
-		var jsonData map[string]interface{}
-		if err := json.Unmarshal(databytes, &jsonData); err != nil {
+	} else if fileExists(os.Getenv("HOME") + configFilePath) {
+		var auth string
+		auth, err = fetchAuthInfo(request.URL.Host, true)
+		if err != nil {
 			return "", err
 		}
-		for _, v := range jsonData {
-			for registry, v1 := range (v).(map[string]interface{}) {
-				for _, credentials := range (v1).(map[string]interface{}) {
-					if registry == request.URL.Host {
-						auth = credentials.(string)
-					}
-				}
-			}
-		}
-		authHeaderVal := "Basic " + auth
+		authHeaderVal = "Basic " + auth
 		response, err = client.doRequest(request, jsonAcceptHeader, authHeaderVal, tokenFetchURL)
 		if err != nil {
 			return "", err
@@ -142,13 +181,13 @@ func (client *orasSourceClient) fetchToken(request *source.Request, path string)
 			return "", err
 		}
 	}
-	token, err := ioutil.ReadAll(response.Body)
+	token, err := io.ReadAll(response.Body)
 	if err != nil {
 		return "", err
 	}
 	defer response.Body.Close()
 	var tokenData map[string]interface{}
-	if err := json.Unmarshal(token, &tokenData); err != nil {
+	if err = json.Unmarshal(token, &tokenData); err != nil {
 		return "", err
 	}
 	logger.Info(fmt.Sprintf("fetching token for %s  successful", request.URL))

pkg/source/source_client_builder.go

@@ -18,13 +18,17 @@ package source
 
 import (
 	"fmt"
+	"net/url"
 
 	"gopkg.in/yaml.v3"
+
+	commonv1 "d7y.io/api/pkg/apis/common/v1"
 )
 
 var (
 	resourceClientBuilder = map[string]ResourceClientBuilder{}
 	resourceClientOptions = map[string]interface{}{}
+	resourceAuthInjector  = map[string]AuthInfoInjector{}
 )
 
 // ResourceClientBuilder is used to build resource client with custom option
@@ -33,12 +37,35 @@ type ResourceClientBuilder interface {
 	Build(optionYaml []byte) (resourceClient ResourceClient, adaptor RequestAdapter, hooks []Hook, err error)
 }
 
+// AuthInfoInjector will inject auth information for target url and metadata, eg: fetch docker config for different users
+type AuthInfoInjector interface {
+	Inject(_url *url.URL, urlMeta *commonv1.UrlMeta) error
+}
+
+// RegisterOption is used for extra options when registering, like mark target scheme protocol should inject auth information
+type RegisterOption func(scheme string)
+
 // RegisterBuilder register ResourceClientBuilder into global resourceClientBuilder, the InitSourceClients will use it.
-func RegisterBuilder(scheme string, builder ResourceClientBuilder) {
+func RegisterBuilder(scheme string, builder ResourceClientBuilder, opts ...RegisterOption) {
 	if _, ok := resourceClientBuilder[scheme]; ok {
 		panic(fmt.Sprintf("duplicate ResourceClientBuilder: %s", scheme))
 	}
 	resourceClientBuilder[scheme] = builder
+
+	for _, opt := range opts {
+		opt(scheme)
+	}
+}
+
+func WithAuthInfoInjector(inj AuthInfoInjector) RegisterOption {
+	return func(scheme string) {
+		resourceAuthInjector[scheme] = inj
+	}
+}
+
+func ShouldInjectAuthInfo(scheme string) (AuthInfoInjector, bool) {
+	inj, ok := resourceAuthInjector[scheme]
+	return inj, ok
 }
 
 func UnRegisterBuilder(scheme string) {
@@ -88,3 +115,16 @@ func NewPlainResourceClientBuilder(
 	build func(optionYaml []byte) (resourceClient ResourceClient, adaptor RequestAdapter, hooks []Hook, err error)) ResourceClientBuilder {
 	return &plainResourceClientBuilder{build: build}
 }
+
+type plainAuthInfoInjector struct {
+	inject func(url *url.URL, urlMeta *commonv1.UrlMeta) error
+}
+
+func (a *plainAuthInfoInjector) Inject(_url *url.URL, urlMeta *commonv1.UrlMeta) error {
+	return a.inject(_url, urlMeta)
+}
+
+func NewPlainAuthInfoInjector(
+	inject func(url *url.URL, urlMeta *commonv1.UrlMeta) error) AuthInfoInjector {
+	return &plainAuthInfoInjector{inject: inject}
+}