diff options
author | Igor Drozdov <idrozdov@gitlab.com> | 2023-01-30 16:00:48 +0100 |
---|---|---|
committer | Igor Drozdov <idrozdov@gitlab.com> | 2023-01-30 19:39:59 +0100 |
commit | 302b7c2ba7b64724cfdf3d46217256b8b738ae3a (patch) | |
tree | 4fa95fc8ea62b3019106216294666342db06784a | |
parent | 51eab44edafd0c097e82c1a74fd379cae4869a42 (diff) | |
download | gitlab-shell-tmp-geo-push-poc.tar.gz |
Poc: Send Geo requests as streamed multipart requeststmp-geo-push-poc
-rw-r--r-- | client/gitlabnet.go | 92 | ||||
-rw-r--r-- | internal/command/shared/customaction/customaction.go | 89 |
2 files changed, 108 insertions, 73 deletions
diff --git a/client/gitlabnet.go b/client/gitlabnet.go index 38adf2a..43502ca 100644 --- a/client/gitlabnet.go +++ b/client/gitlabnet.go @@ -150,21 +150,10 @@ func (c *GitlabNetClient) Post(ctx context.Context, path string, data interface{ return c.DoRequest(ctx, http.MethodPost, normalizePath(path), data) } -func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, data interface{}) (*http.Response, error) { - request, err := newRequest(ctx, method, c.httpClient.Host, path, data) - if err != nil { - return nil, err - } - - retryableRequest, err := newRetryableRequest(ctx, method, c.httpClient.Host, path, data) - if err != nil { - return nil, err - } - +func (c *GitlabNetClient) prepareRequest(request *http.Request) error { user, password := c.user, c.password if user != "" && password != "" { request.SetBasicAuth(user, password) - retryableRequest.SetBasicAuth(user, password) } claims := jwt.RegisteredClaims{ @@ -175,44 +164,33 @@ func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, da secretBytes := []byte(strings.TrimSpace(c.secret)) tokenString, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString(secretBytes) if err != nil { - return nil, err + return err } request.Header.Set(apiSecretHeaderName, tokenString) - retryableRequest.Header.Set(apiSecretHeaderName, tokenString) - originalRemoteIP, ok := ctx.Value(OriginalRemoteIPContextKey{}).(string) + originalRemoteIP, ok := request.Context().Value(OriginalRemoteIPContextKey{}).(string) if ok { request.Header.Add("X-Forwarded-For", originalRemoteIP) - retryableRequest.Header.Add("X-Forwarded-For", originalRemoteIP) } request.Header.Add("Content-Type", "application/json") - retryableRequest.Header.Add("Content-Type", "application/json") request.Header.Add("User-Agent", c.userAgent) - retryableRequest.Header.Add("User-Agent", c.userAgent) request.Close = true - retryableRequest.Close = true - start := time.Now() + return nil +} - var response *http.Response - var respErr error - if c.httpClient.HTTPClient != nil { - response, respErr = c.httpClient.HTTPClient.Do(request) - } - if os.Getenv("FF_GITLAB_SHELL_RETRYABLE_HTTP") == "1" && c.httpClient.RetryableHTTP != nil { - response, respErr = c.httpClient.RetryableHTTP.Do(retryableRequest) - } +func processResult(request *http.Request, response *http.Response, start time.Time, respErr error) error { fields := log.Fields{ - "method": method, + "method": request.Method, "url": request.URL.String(), "duration_ms": time.Since(start) / time.Millisecond, } - logger := log.WithContextFields(ctx, fields) + logger := log.WithContextFields(request.Context(), fields) if respErr != nil { logger.WithError(respErr).Error("Internal API unreachable") - return nil, &ApiError{"Internal API unreachable"} + return &ApiError{"Internal API unreachable"} } if response != nil { @@ -220,7 +198,7 @@ func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, da } if err := parseError(response); err != nil { logger.WithError(err).Error("Internal API error") - return nil, err + return err } if response.ContentLength >= 0 { @@ -229,5 +207,55 @@ func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, da logger.Info("Finished HTTP request") + return nil +} + +func (c *GitlabNetClient) AppendPath(path string) string { + return appendPath(c.httpClient.Host, path) +} + +func (c *GitlabNetClient) DoRawRequest(request *http.Request) (*http.Response, error) { + c.prepareRequest(request) + + start := time.Now() + + response, respErr := c.httpClient.HTTPClient.Do(request) + + if err := processResult(request, response, start, respErr); err != nil { + return nil, err + } + return response, nil } + +func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, data interface{}) (*http.Response, error) { + request, err := newRequest(ctx, method, c.httpClient.Host, path, data) + if err != nil { + return nil, err + } + + retryableRequest, err := newRetryableRequest(ctx, method, c.httpClient.Host, path, data) + if err != nil { + return nil, err + } + + c.prepareRequest(request) + c.prepareRequest(retryableRequest.Request) + + start := time.Now() + + var response *http.Response + var respErr error + if c.httpClient.HTTPClient != nil { + response, respErr = c.httpClient.HTTPClient.Do(request) + } + if os.Getenv("FF_GITLAB_SHELL_RETRYABLE_HTTP") == "1" && c.httpClient.RetryableHTTP != nil { + response, respErr = c.httpClient.RetryableHTTP.Do(retryableRequest) + } + + if err := processResult(request, response, start, respErr); err != nil { + return nil, err + } + + return response, err +} diff --git a/internal/command/shared/customaction/customaction.go b/internal/command/shared/customaction/customaction.go index c12d685..f9341ad 100644 --- a/internal/command/shared/customaction/customaction.go +++ b/internal/command/shared/customaction/customaction.go @@ -1,11 +1,11 @@ package customaction import ( - "bytes" "context" "errors" "io" "net/http" + "mime/multipart" "gitlab.com/gitlab-org/labkit/log" @@ -18,14 +18,12 @@ import ( ) type Request struct { - SecretToken []byte `json:"secret_token"` Data accessverifier.CustomPayloadData `json:"data"` - Output []byte `json:"output"` + Output io.Reader } type Response struct { Result []byte `json:"result"` - Message string `json:"message"` } type Command struct { @@ -54,7 +52,7 @@ func (c *Command) processApiEndpoints(ctx context.Context, response *accessverif data := response.Payload.Data request := &Request{Data: data} - request.Data.UserId = response.Who + request.Data.UserId = response.Who for _, endpoint := range data.ApiEndpoints { ctxlog := log.WithContextFields(ctx, log.Fields{ @@ -64,64 +62,82 @@ func (c *Command) processApiEndpoints(ctx context.Context, response *accessverif ctxlog.Info("customaction: processApiEndpoints: Performing custom action") - response, err := c.performRequest(ctx, client, endpoint, request) + httpRequest, err := c.prepareRequest(ctx, client.AppendPath(endpoint), request) if err != nil { return err } - // Print to os.Stdout the result contained in the response - // - if err = c.displayResult(response.Result); err != nil { + if err := c.performRequest(ctx, client, httpRequest); err != nil { return err } // In the context of the git push sequence of events, it's necessary to read // stdin in order to capture output to pass onto subsequent commands - // - var output []byte - if c.EOFSent { - output, err = c.readFromStdin() - if err != nil { - return err - } + var w *io.PipeWriter + request.Output, w = io.Pipe() + + go c.readFromStdin(w) } else { - output = c.readFromStdinNoEOF() + // output = c.readFromStdinNoEOF() } ctxlog.WithFields(log.Fields{ "eof_sent": c.EOFSent, - "stdin_bytes": len(output), + // "stdin_bytes": len(output), }).Debug("customaction: processApiEndpoints: stdin buffered") - - request.Output = output } return nil } -func (c *Command) performRequest(ctx context.Context, client *client.GitlabNetClient, endpoint string, request *Request) (*Response, error) { - response, err := client.DoRequest(ctx, http.MethodPost, endpoint, request) +func (c *Command) prepareRequest(ctx context.Context, endpoint string, request *Request) (*http.Request, error) { + body, pipeWriter := io.Pipe() + writer := multipart.NewWriter(pipeWriter) + + go func() { + writer.WriteField("data[gl_id]", request.Data.UserId) + writer.WriteField("data[primary_repo]", request.Data.PrimaryRepo) + + if request.Output != nil { + // Ignore errors, but may want to log them in a channel + binaryPart, _ := writer.CreateFormFile("output", "git-receive-pack") + io.Copy(binaryPart, request.Output) + } + + writer.Close() + pipeWriter.Close() + }() + + httpRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, body) if err != nil { return nil, err } + httpRequest.Header.Set("Content-Type", writer.FormDataContentType()) + + return httpRequest, nil +} + +func (c *Command) performRequest(ctx context.Context, client *client.GitlabNetClient, request *http.Request) error { + response, err := client.DoRawRequest(request) + if err != nil { + return err + } defer response.Body.Close() - cr := &Response{} - if err := gitlabnet.ParseJSON(response, cr); err != nil { - return nil, err + if _, err := io.Copy(c.ReadWriter.Out, response.Body); err != nil { + return err } - return cr, nil + return nil } -func (c *Command) readFromStdin() ([]byte, error) { - var output []byte +func (c *Command) readFromStdin(w *io.PipeWriter) { var needsPackData bool scanner := pktline.NewScanner(c.ReadWriter.In) for scanner.Scan() { line := scanner.Bytes() - output = append(output, line...) + w.Write(line) if pktline.IsFlush(line) { break @@ -133,14 +149,10 @@ func (c *Command) readFromStdin() ([]byte, error) { } if needsPackData { - packData := new(bytes.Buffer) - _, err := io.Copy(packData, c.ReadWriter.In) - - output = append(output, packData.Bytes()...) - return output, err - } else { - return output, nil + io.Copy(w, c.ReadWriter.In) } + + w.Close() } func (c *Command) readFromStdinNoEOF() []byte { @@ -158,8 +170,3 @@ func (c *Command) readFromStdinNoEOF() []byte { return output } - -func (c *Command) displayResult(result []byte) error { - _, err := io.Copy(c.ReadWriter.Out, bytes.NewReader(result)) - return err -} |