Skip to content

Commit 55530da

Browse files
fix: fix for file uploads to octet stream and form encoding endpoints
1 parent dc0fd28 commit 55530da

File tree

2 files changed

+117
-67
lines changed

2 files changed

+117
-67
lines changed

pkg/cmd/flagoptions.go

Lines changed: 90 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cmd
22

33
import (
4+
"bufio"
45
"bytes"
56
"encoding/base64"
67
"encoding/json"
@@ -33,17 +34,24 @@ const (
3334
ApplicationOctetStream
3435
)
3536

36-
func embedFiles(obj any) (any, error) {
37+
type FileEmbedStyle int
38+
39+
const (
40+
EmbedText FileEmbedStyle = iota
41+
EmbedIOReader
42+
)
43+
44+
func embedFiles(obj any, embedStyle FileEmbedStyle) (any, error) {
3745
v := reflect.ValueOf(obj)
38-
result, err := embedFilesValue(v)
46+
result, err := embedFilesValue(v, embedStyle)
3947
if err != nil {
4048
return nil, err
4149
}
4250
return result.Interface(), nil
4351
}
4452

4553
// Replace "@file.txt" with the file's contents inside a value
46-
func embedFilesValue(v reflect.Value) (reflect.Value, error) {
54+
func embedFilesValue(v reflect.Value, embedStyle FileEmbedStyle) (reflect.Value, error) {
4755
// Unwrap interface values to get the concrete type
4856
if v.Kind() == reflect.Interface {
4957
if v.IsNil() {
@@ -57,12 +65,14 @@ func embedFilesValue(v reflect.Value) (reflect.Value, error) {
5765
if v.Len() == 0 {
5866
return v, nil
5967
}
60-
result := reflect.MakeMap(v.Type())
68+
// Always create map[string]any to handle potential type changes when embedding files
69+
result := reflect.MakeMap(reflect.TypeOf(map[string]any{}))
70+
6171
iter := v.MapRange()
6272
for iter.Next() {
6373
key := iter.Key()
6474
val := iter.Value()
65-
newVal, err := embedFilesValue(val)
75+
newVal, err := embedFilesValue(val, embedStyle)
6676
if err != nil {
6777
return reflect.Value{}, err
6878
}
@@ -74,9 +84,10 @@ func embedFilesValue(v reflect.Value) (reflect.Value, error) {
7484
if v.Len() == 0 {
7585
return v, nil
7686
}
77-
result := reflect.MakeSlice(v.Type(), v.Len(), v.Len())
87+
// Use `[]any` to allow for types to change when embedding files
88+
result := reflect.MakeSlice(reflect.TypeOf([]any{}), v.Len(), v.Len())
7889
for i := 0; i < v.Len(); i++ {
79-
newVal, err := embedFilesValue(v.Index(i))
90+
newVal, err := embedFilesValue(v.Index(i), embedStyle)
8091
if err != nil {
8192
return reflect.Value{}, err
8293
}
@@ -86,51 +97,78 @@ func embedFilesValue(v reflect.Value) (reflect.Value, error) {
8697

8798
case reflect.String:
8899
s := v.String()
89-
90100
if literal, ok := strings.CutPrefix(s, "\\@"); ok {
91101
// Allow for escaped @ signs if you don't want them to be treated as files
92102
return reflect.ValueOf("@" + literal), nil
93-
} else if filename, ok := strings.CutPrefix(s, "@data://"); ok {
94-
// The "@data://" prefix is for files you explicitly want to upload
95-
// as base64-encoded (even if the file itself is plain text)
96-
content, err := os.ReadFile(filename)
97-
if err != nil {
98-
return v, err
99-
}
100-
return reflect.ValueOf(base64.StdEncoding.EncodeToString(content)), nil
101-
} else if filename, ok := strings.CutPrefix(s, "@file://"); ok {
102-
// The "@file://" prefix is for files that you explicitly want to
103-
// upload as a string literal with backslash escapes (not base64
104-
// encoded)
105-
content, err := os.ReadFile(filename)
106-
if err != nil {
107-
return v, err
108-
}
109-
return reflect.ValueOf(string(content)), nil
110-
} else if filename, ok := strings.CutPrefix(s, "@"); ok {
111-
content, err := os.ReadFile(filename)
112-
if err != nil {
113-
// If the string is "@username", it's probably supposed to be a
114-
// string literal and not a file reference. However, if the
115-
// string looks like "@file.txt" or "@/tmp/file", then it's
116-
// probably supposed to be a file.
117-
probablyFile := strings.Contains(filename, ".") || strings.Contains(filename, "/")
118-
if probablyFile {
119-
// Give a useful error message if the user tried to upload a
120-
// file, but the file couldn't be read (e.g. mistyped
121-
// filename or permission error)
103+
}
104+
105+
if embedStyle == EmbedText {
106+
if filename, ok := strings.CutPrefix(s, "@data://"); ok {
107+
// The "@data://" prefix is for files you explicitly want to upload
108+
// as base64-encoded (even if the file itself is plain text)
109+
content, err := os.ReadFile(filename)
110+
if err != nil {
111+
return v, err
112+
}
113+
return reflect.ValueOf(base64.StdEncoding.EncodeToString(content)), nil
114+
} else if filename, ok := strings.CutPrefix(s, "@file://"); ok {
115+
// The "@file://" prefix is for files that you explicitly want to
116+
// upload as a string literal with backslash escapes (not base64
117+
// encoded)
118+
content, err := os.ReadFile(filename)
119+
if err != nil {
122120
return v, err
123121
}
124-
// Fall back to the raw value if the user provided something
125-
// like "@username" that's not intended to be a file.
126-
return v, nil
127-
}
128-
// If the file looks like a plain text UTF8 file format, then use the contents directly.
129-
if isUTF8TextFile(content) {
130122
return reflect.ValueOf(string(content)), nil
123+
} else if filename, ok := strings.CutPrefix(s, "@"); ok {
124+
content, err := os.ReadFile(filename)
125+
if err != nil {
126+
// If the string is "@username", it's probably supposed to be a
127+
// string literal and not a file reference. However, if the
128+
// string looks like "@file.txt" or "@/tmp/file", then it's
129+
// probably supposed to be a file.
130+
probablyFile := strings.Contains(filename, ".") || strings.Contains(filename, "/")
131+
if probablyFile {
132+
// Give a useful error message if the user tried to upload a
133+
// file, but the file couldn't be read (e.g. mistyped
134+
// filename or permission error)
135+
return v, err
136+
}
137+
// Fall back to the raw value if the user provided something
138+
// like "@username" that's not intended to be a file.
139+
return v, nil
140+
}
141+
// If the file looks like a plain text UTF8 file format, then use the contents directly.
142+
if isUTF8TextFile(content) {
143+
return reflect.ValueOf(string(content)), nil
144+
}
145+
// Otherwise it's a binary file, so encode it with base64
146+
return reflect.ValueOf(base64.StdEncoding.EncodeToString(content)), nil
147+
}
148+
} else {
149+
if filename, ok := strings.CutPrefix(s, "@"); ok {
150+
// Behavior is the same for @file, @data://file, and @file://file, except that
151+
// @username will be treated as a literal string if no "username" file exists
152+
expectsFile := true
153+
if withoutPrefix, ok := strings.CutPrefix(filename, "data://"); ok {
154+
filename = withoutPrefix
155+
} else if withoutPrefix, ok := strings.CutPrefix(filename, "file://"); ok {
156+
filename = withoutPrefix
157+
} else {
158+
expectsFile = strings.Contains(filename, ".") || strings.Contains(filename, "/")
159+
}
160+
161+
file, err := os.Open(filename)
162+
if err != nil {
163+
if !expectsFile {
164+
// For strings that start with "@" and don't look like a filename, return the string
165+
return v, nil
166+
}
167+
return v, err
168+
}
169+
reader := bufio.NewReader(file)
170+
return reflect.ValueOf(reader), nil
131171
}
132-
// Otherwise it's a binary file, so encode it with base64
133-
return reflect.ValueOf(base64.StdEncoding.EncodeToString(content)), nil
134172
}
135173
return v, nil
136174

@@ -205,16 +243,20 @@ func flagOptions(
205243
}
206244

207245
// Embed files passed as "@file.jpg" in the request body, headers, and query:
208-
bodyData, err := embedFiles(bodyData)
246+
embedStyle := EmbedText
247+
if bodyType == ApplicationOctetStream || bodyType == MultipartFormEncoded {
248+
embedStyle = EmbedIOReader
249+
}
250+
bodyData, err := embedFiles(bodyData, embedStyle)
209251
if err != nil {
210252
return nil, err
211253
}
212-
if headersWithFiles, err := embedFiles(flagContents.Headers); err != nil {
254+
if headersWithFiles, err := embedFiles(flagContents.Headers, EmbedText); err != nil {
213255
return nil, err
214256
} else {
215257
flagContents.Headers = headersWithFiles.(map[string]any)
216258
}
217-
if queriesWithFiles, err := embedFiles(flagContents.Queries); err != nil {
259+
if queriesWithFiles, err := embedFiles(flagContents.Queries, EmbedText); err != nil {
218260
return nil, err
219261
} else {
220262
flagContents.Queries = queriesWithFiles.(map[string]any)

pkg/cmd/flagoptions_test.go

Lines changed: 27 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -68,11 +68,11 @@ func TestEmbedFiles(t *testing.T) {
6868
},
6969
{
7070
name: "map[string]string with file references",
71-
input: map[string]string{
71+
input: map[string]any{
7272
"config": "@" + filepath.Join(tmpDir, "config.txt"),
7373
"name": "test",
7474
},
75-
want: map[string]string{
75+
want: map[string]any{
7676
"config": configContent,
7777
"name": "test",
7878
},
@@ -96,11 +96,11 @@ func TestEmbedFiles(t *testing.T) {
9696
},
9797
{
9898
name: "[]string with file references",
99-
input: []string{
99+
input: []any{
100100
"@" + filepath.Join(tmpDir, "config.txt"),
101101
"normal string",
102102
},
103-
want: []string{
103+
want: []any{
104104
configContent,
105105
"normal string",
106106
},
@@ -112,7 +112,7 @@ func TestEmbedFiles(t *testing.T) {
112112
"outer": map[string]any{
113113
"inner": []any{
114114
"@" + filepath.Join(tmpDir, "config.txt"),
115-
map[string]string{
115+
map[string]any{
116116
"data": "@" + filepath.Join(tmpDir, "data.json"),
117117
},
118118
},
@@ -122,7 +122,7 @@ func TestEmbedFiles(t *testing.T) {
122122
"outer": map[string]any{
123123
"inner": []any{
124124
configContent,
125-
map[string]string{
125+
map[string]any{
126126
"data": dataContent,
127127
},
128128
},
@@ -132,53 +132,53 @@ func TestEmbedFiles(t *testing.T) {
132132
},
133133
{
134134
name: "base64 encoding",
135-
input: map[string]string{
135+
input: map[string]any{
136136
"encoded": "@data://" + filepath.Join(tmpDir, "config.txt"),
137137
"image": "@" + filepath.Join(tmpDir, "image.jpg"),
138138
},
139-
want: map[string]string{
139+
want: map[string]any{
140140
"encoded": base64.StdEncoding.EncodeToString([]byte(configContent)),
141141
"image": base64.StdEncoding.EncodeToString(jpegHeader),
142142
},
143143
wantErr: false,
144144
},
145145
{
146146
name: "non-existent file with @ prefix",
147-
input: map[string]string{
147+
input: map[string]any{
148148
"missing": "@file.txt",
149149
},
150150
want: nil,
151151
wantErr: true,
152152
},
153153
{
154154
name: "non-file-like thing with @ prefix",
155-
input: map[string]string{
155+
input: map[string]any{
156156
"username": "@user",
157157
"favorite_symbol": "@",
158158
},
159-
want: map[string]string{
159+
want: map[string]any{
160160
"username": "@user",
161161
"favorite_symbol": "@",
162162
},
163163
wantErr: false,
164164
},
165165
{
166166
name: "non-existent file with @file:// prefix (error)",
167-
input: map[string]string{
167+
input: map[string]any{
168168
"missing": "@file:///nonexistent/file.txt",
169169
},
170170
want: nil,
171171
wantErr: true,
172172
},
173173
{
174174
name: "escaping",
175-
input: map[string]string{
175+
input: map[string]any{
176176
"simple": "\\@file.txt",
177177
"file": "\\@file://file.txt",
178178
"data": "\\@data://file.txt",
179179
"keep_escape": "user\\@example.com",
180180
},
181-
want: map[string]string{
181+
want: map[string]any{
182182
"simple": "@file.txt",
183183
"file": "@file://file.txt",
184184
"data": "@data://file.txt",
@@ -207,24 +207,32 @@ func TestEmbedFiles(t *testing.T) {
207207
wantErr: false,
208208
},
209209
{
210-
name: "[]int unchanged",
210+
name: "[]int values unchanged",
211211
input: []int{1, 2, 3, 4, 5},
212-
want: []int{1, 2, 3, 4, 5},
212+
want: []any{1, 2, 3, 4, 5},
213213
wantErr: false,
214214
},
215215
}
216216

217217
for _, tt := range tests {
218-
t.Run(tt.name, func(t *testing.T) {
219-
got, err := embedFiles(tt.input)
220-
218+
t.Run(tt.name+" text", func(t *testing.T) {
219+
got, err := embedFiles(tt.input, EmbedText)
221220
if tt.wantErr {
222221
assert.Error(t, err)
223222
} else {
224223
require.NoError(t, err)
225224
assert.Equal(t, tt.want, got)
226225
}
227226
})
227+
228+
t.Run(tt.name+" io.Reader", func(t *testing.T) {
229+
_, err := embedFiles(tt.input, EmbedIOReader)
230+
if tt.wantErr {
231+
assert.Error(t, err)
232+
} else {
233+
require.NoError(t, err)
234+
}
235+
})
228236
}
229237
}
230238

0 commit comments

Comments
 (0)