Skip to content

Commit 828c770

Browse files
ShocOneShocOne
committed
Refactor header parsing functions in response package
1 parent f55f77f commit 828c770

File tree

2 files changed

+106
-103
lines changed

2 files changed

+106
-103
lines changed

‎response/parse.go‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,31 @@ package response
33

44
import "strings"
55

6-
// ParseContentTypeHeader parses the Content-Type header and extracts the MIME type and parameters.
6+
// ParseContentTypeHeader parses the Content-Type header and returns the MIME type and any parameters.
77
func ParseContentTypeHeader(header string) (string, map[string]string) {
8-
parts := strings.Split(header, ";")
9-
mimeType := strings.TrimSpace(parts[0])
8+
return parseHeader(header)
9+
}
10+
11+
// ParseContentDisposition parses the Content-Disposition header and returns the type and any parameters.
12+
func ParseContentDisposition(header string) (string, map[string]string) {
13+
return parseHeader(header)
14+
}
15+
16+
// parseHeader generalizes the parsing of headers like Content-Type and Content-Disposition.
17+
// It extracts the main value (e.g., MIME type for Content-Type) and any parameters (like charset).
18+
func parseHeader(header string) (string, map[string]string) {
19+
parts := strings.SplitN(header, ";", 2) // Split into main value and parameters
20+
mainValue := strings.TrimSpace(parts[0])
21+
1022
params := make(map[string]string)
11-
for _, part := range parts[1:] {
12-
kv := strings.SplitN(part, "=", 2)
13-
if len(kv) == 2 {
14-
params[strings.TrimSpace(kv[0])] = strings.TrimSpace(kv[1])
23+
if len(parts) > 1 { // Check if there are parameters
24+
for _, part := range strings.Split(parts[1], ";") {
25+
kv := strings.SplitN(part, "=", 2)
26+
if len(kv) == 2 {
27+
params[strings.TrimSpace(kv[0])] = strings.Trim(strings.TrimSpace(kv[1]), "\"")
28+
}
1529
}
1630
}
17-
return mimeType, params
31+
32+
return mainValue, params
1833
}

‎response/success.go‎

Lines changed: 83 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ package response
66
import (
77
"encoding/json"
88
"encoding/xml"
9+
"errors"
910
"fmt"
1011
"io"
1112
"net/http"
@@ -15,32 +16,39 @@ import (
1516
"go.uber.org/zap"
1617
)
1718

18-
// HandleAPISuccessResponse handles the HTTP success response from an API and unmarshals the response body into the provided output struct.
19+
// Refactored contentHandler to accept io.Reader instead of []byte for streaming support.
20+
type contentHandler func(io.Reader, interface{}, logger.Logger, string) error
21+
22+
// Updated handlers map to use the new contentHandler signature.
23+
var handlers = map[string]contentHandler{
24+
"application/json": unmarshalJSON,
25+
"application/xml": unmarshalXML,
26+
"text/xml": unmarshalXML,
27+
}
28+
29+
// HandleAPISuccessResponse reads the response body and unmarshals it based on the content type.
1930
func HandleAPISuccessResponse(resp *http.Response, out interface{}, log logger.Logger) error {
20-
// Special handling for DELETE requests
2131
if resp.Request.Method == "DELETE" {
2232
return handleDeleteRequest(resp, log)
2333
}
2434

25-
// Read the response body
26-
bodyBytes, err := readResponseBody(resp, log)
27-
if err != nil {
28-
return err
29-
}
30-
31-
// Log the raw response details for debugging
32-
logResponseDetails(resp, bodyBytes, log)
35+
// No need to read the entire body into memory, pass resp.Body directly.
36+
logResponseDetails(resp, nil, log) // Updated to handle nil bodyBytes.
3337

34-
// Unmarshal the response based on content type
35-
contentType := resp.Header.Get("Content-Type")
36-
37-
// Check for binary data handling
38+
mimeType, _ := ParseContentTypeHeader(resp.Header.Get("Content-Type"))
3839
contentDisposition := resp.Header.Get("Content-Disposition")
39-
if err := handleBinaryData(contentType, contentDisposition, bodyBytes, log, out); err != nil {
40-
return err
41-
}
4240

43-
return unmarshalResponse(contentType, bodyBytes, log, out)
41+
if handler, ok := handlers[mimeType]; ok {
42+
// Pass resp.Body directly to the handler for streaming.
43+
return handler(resp.Body, out, log, mimeType)
44+
} else if isBinaryData(mimeType, contentDisposition) {
45+
// For binary data, we still need to handle the body directly.
46+
return handleBinaryData(resp.Body, log, out, mimeType, contentDisposition)
47+
} else {
48+
errMsg := fmt.Sprintf("unexpected MIME type: %s", mimeType)
49+
log.Error("Unmarshal error", zap.String("content type", mimeType), zap.Error(errors.New(errMsg)))
50+
return errors.New(errMsg)
51+
}
4452
}
4553

4654
// handleDeleteRequest handles the special case for DELETE requests, where a successful response might not contain a body.
@@ -52,98 +60,78 @@ func handleDeleteRequest(resp *http.Response, log logger.Logger) error {
5260
return log.Error("DELETE request failed", zap.String("URL", resp.Request.URL.String()), zap.Int("Status Code", resp.StatusCode))
5361
}
5462

55-
// readResponseBody reads and returns the body of an HTTP response. It logs an error if reading fails.
56-
func readResponseBody(resp *http.Response, log logger.Logger) ([]byte, error) {
57-
// Read the response body
58-
bodyBytes, err := io.ReadAll(resp.Body)
59-
if err != nil {
60-
log.Error("Failed reading response body", zap.Error(err))
61-
return nil, err
62-
}
63-
return bodyBytes, nil
64-
}
65-
66-
// logResponseDetails logs the raw HTTP response body and headers for debugging purposes.
63+
// Adjusted logResponseDetails to handle a potential nil bodyBytes.
6764
func logResponseDetails(resp *http.Response, bodyBytes []byte, log logger.Logger) {
68-
// Log the response body as a string
69-
log.Debug("Raw HTTP Response", zap.String("Body", string(bodyBytes)))
70-
// Log the response headers
65+
// Conditional logging if bodyBytes is not nil.
66+
if bodyBytes != nil {
67+
log.Debug("Raw HTTP Response", zap.String("Body", string(bodyBytes)))
68+
}
69+
// Logging headers remains unchanged.
7170
log.Debug("HTTP Response Headers", zap.Any("Headers", resp.Header))
7271
}
7372

73+
// unmarshalJSON unmarshals JSON content from an io.Reader into the provided output structure.
74+
func unmarshalJSON(reader io.Reader, out interface{}, log logger.Logger, mimeType string) error {
75+
decoder := json.NewDecoder(reader)
76+
if err := decoder.Decode(out); err != nil {
77+
log.Error("JSON Unmarshal error", zap.Error(err))
78+
return err
79+
}
80+
log.Info("Successfully unmarshalled JSON response", zap.String("content type", mimeType))
81+
return nil
82+
}
7483

75-
// handleBinaryData checks if the response should be treated as binary data based on the Content-Type or Content-Disposition headers.
76-
// It supports writing the response body to an output parameter of type *[]byte or io.Writer.
77-
func handleBinaryData(contentType, contentDisposition string, bodyBytes []byte, log logger.Logger, out interface{}) error {
78-
// Check if response is binary data either by Content-Type or Content-Disposition
79-
if strings.Contains(contentType, "application/octet-stream") || strings.HasPrefix(contentDisposition, "attachment") {
80-
switch out := out.(type) {
81-
case *[]byte:
82-
*out = bodyBytes // Assign the response body to 'out'
83-
log.Debug("Handled binary data as []byte",
84-
zap.String("Content-Type", contentType),
85-
zap.String("Content-Disposition", contentDisposition),
86-
)
87-
88-
case io.Writer:
89-
_, err := out.Write(bodyBytes) // Write the response body to 'out'
90-
if err != nil {
91-
errMsg := "failed to write binary data to io.Writer"
92-
log.Error("Binary data writing error",
93-
zap.String("error", errMsg),
94-
zap.String("Content-Type", contentType),
95-
zap.String("Content-Disposition", contentDisposition),
96-
)
97-
return fmt.Errorf(errMsg)
98-
}
99-
log.Debug("Handled binary data as io.Writer",
100-
zap.String("Content-Type", contentType),
101-
zap.String("Content-Disposition", contentDisposition),
102-
)
103-
104-
default:
105-
errMsg := "output parameter is not a *[]byte or io.Writer for binary data"
106-
log.Error("Binary data handling error",
107-
zap.String("error", errMsg),
108-
zap.String("Content-Type", contentType),
109-
zap.String("Content-Disposition", contentDisposition),
110-
)
111-
return fmt.Errorf(errMsg)
112-
}
113-
return nil
84+
// unmarshalXML unmarshals XML content from an io.Reader into the provided output structure.
85+
func unmarshalXML(reader io.Reader, out interface{}, log logger.Logger, mimeType string) error {
86+
decoder := xml.NewDecoder(reader)
87+
if err := decoder.Decode(out); err != nil {
88+
log.Error("XML Unmarshal error", zap.Error(err))
89+
return err
11490
}
115-
return nil // If not binary data, no action needed
91+
log.Info("Successfully unmarshalled XML response", zap.String("content type", mimeType))
92+
return nil
11693
}
11794

118-
// unmarshalResponse unmarshals the response body into the provided output structure based on the MIME
119-
// type extracted from the Content-Type header.
120-
func unmarshalResponse(contentTypeHeader string, bodyBytes []byte, log logger.Logger, out interface{}) error {
121-
// Extract MIME type from Content-Type header
122-
mimeType, _ := ParseContentTypeHeader(contentTypeHeader)
123-
124-
// Determine the MIME type and unmarshal accordingly
125-
switch {
126-
case strings.Contains(mimeType, "application/json"):
127-
// Unmarshal JSON content
128-
if err := json.Unmarshal(bodyBytes, out); err != nil {
129-
log.Error("JSON Unmarshal error", zap.Error(err))
95+
// isBinaryData checks if the MIME type or Content-Disposition indicates binary data.
96+
func isBinaryData(contentType, contentDisposition string) bool {
97+
return strings.Contains(contentType, "application/octet-stream") || strings.HasPrefix(contentDisposition, "attachment")
98+
}
99+
100+
// handleBinaryData reads binary data from an io.Reader and stores it in *[]byte or streams it to an io.Writer.
101+
func handleBinaryData(reader io.Reader, log logger.Logger, out interface{}, mimeType, contentDisposition string) error {
102+
// Check if the output interface is either *[]byte or io.Writer
103+
switch out := out.(type) {
104+
case *[]byte:
105+
// Read all data from reader and store it in *[]byte
106+
data, err := io.ReadAll(reader)
107+
if err != nil {
108+
log.Error("Failed to read binary data", zap.Error(err))
130109
return err
131110
}
132-
log.Info("Successfully unmarshalled JSON response", zap.String("content type", mimeType))
111+
*out = data
133112

134-
case strings.Contains(mimeType, "application/xml") || strings.Contains(mimeType, "text/xml"):
135-
// Unmarshal XML content
136-
if err := xml.Unmarshal(bodyBytes, out); err != nil {
137-
log.Error("XML Unmarshal error", zap.Error(err))
113+
case io.Writer:
114+
// Stream data directly to the io.Writer
115+
_, err := io.Copy(out, reader)
116+
if err != nil {
117+
log.Error("Failed to stream binary data to io.Writer", zap.Error(err))
138118
return err
139119
}
140-
log.Info("Successfully unmarshalled XML response", zap.String("content type", mimeType))
141120

142121
default:
143-
// Log and return an error for unexpected MIME types
144-
errMsg := fmt.Sprintf("unexpected MIME type: %s", mimeType)
145-
log.Error("Unmarshal error", zap.String("content type", mimeType), zap.Error(fmt.Errorf(errMsg)))
146-
return fmt.Errorf(errMsg)
122+
errMsg := "output parameter is not suitable for binary data (*[]byte or io.Writer)"
123+
log.Error(errMsg, zap.String("Content-Type", mimeType))
124+
return errors.New(errMsg)
147125
}
126+
127+
// Handle Content-Disposition if present
128+
if contentDisposition != "" {
129+
_, params := ParseContentDisposition(contentDisposition)
130+
if filename, ok := params["filename"]; ok {
131+
log.Debug("Extracted filename from Content-Disposition", zap.String("filename", filename))
132+
// Additional processing for the filename can be done here if needed
133+
}
134+
}
135+
148136
return nil
149137
}

0 commit comments

Comments
 (0)