Skip to content

Commit 67ea8ec

Browse files
caffixcaffix
committed
improvements to the engine API client
1 parent 8063480 commit 67ea8ec

File tree

3 files changed

+56
-61
lines changed

3 files changed

+56
-61
lines changed

‎engine/api/client/v1/client.go‎

Lines changed: 52 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,10 @@
55
package v1
66

77
import (
8-
"bytes"
98
"context"
109
"encoding/json"
1110
"fmt"
11+
"io"
1212
"net/http"
1313
"net/url"
1414
"os"
@@ -20,14 +20,15 @@ import (
2020
"github.com/owasp-amass/amass/v5/config"
2121
apiclient "github.com/owasp-amass/amass/v5/engine/api/client"
2222
et "github.com/owasp-amass/amass/v5/engine/types"
23+
amasshttp "github.com/owasp-amass/amass/v5/internal/net/http"
2324
oam "github.com/owasp-amass/open-asset-model"
2425
)
2526

26-
const MaxBulkItems = 5000
27+
const MaxBulkItems = 1000
2728

2829
type Client struct {
2930
base string
30-
httpClient http.Client
31+
httpClient *http.Client
3132
wsClient *websocket.Conn
3233
done chan struct{}
3334
}
@@ -60,7 +61,7 @@ type BulkAddAssetsResponse struct {
6061
func NewClient(url string) (*Client, error) {
6162
return &Client{
6263
base: url + "/api/v1",
63-
httpClient: http.Client{},
64+
httpClient: amasshttp.DefaultClient,
6465
done: make(chan struct{}),
6566
}, nil
6667
}
@@ -72,14 +73,10 @@ func (c *Client) Close() {
7273

7374
// HealthCheck returns true when the client was able to reach the server.
7475
func (c *Client) HealthCheck(ctx context.Context) bool {
75-
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/health", nil)
76-
77-
resp, err := c.httpClient.Do(req)
76+
resp, err := amasshttp.RequestWebPage(ctx, c.httpClient, &amasshttp.Request{URL: c.base + "/health"})
7877
if err != nil {
7978
return false
8079
}
81-
defer func() { _ = resp.Body.Close() }()
82-
8380
return resp.StatusCode == http.StatusOK
8481
}
8582

@@ -90,25 +87,26 @@ func (c *Client) CreateSession(ctx context.Context, config *config.Config) (uuid
9087
return uuid.UUID{}, err
9188
}
9289

93-
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, c.base+"/sessions", bytes.NewReader(raw))
94-
req.Header.Set("Content-Type", "application/json")
95-
96-
resp, err := c.httpClient.Do(req)
90+
resp, err := amasshttp.RequestWebPage(ctx, c.httpClient, &amasshttp.Request{
91+
Method: http.MethodPost,
92+
Body: string(raw),
93+
URL: c.base + "/sessions",
94+
Header: amasshttp.Header{"Content-Type": []string{"application/json"}},
95+
})
9796
if err != nil {
9897
return uuid.UUID{}, err
9998
}
100-
defer func() { _ = resp.Body.Close() }()
10199

102100
if resp.StatusCode != http.StatusCreated {
103-
msg, err := readJSONError(resp)
101+
msg, err := readJSONError(resp.Body)
104102
if err != nil {
105103
return uuid.UUID{}, fmt.Errorf("createSession: status=%s", resp.Status)
106104
}
107105
return uuid.UUID{}, fmt.Errorf("createSession: status=%s error=%s", resp.Status, msg)
108106
}
109107

110108
var out CreateSessionResponse
111-
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
109+
if err := json.Unmarshal([]byte(resp.Body), &out); err != nil {
112110
return uuid.UUID{}, err
113111
}
114112

@@ -117,24 +115,21 @@ func (c *Client) CreateSession(ctx context.Context, config *config.Config) (uuid
117115

118116
// Lists the active session and associated tokens on the server.
119117
func (c *Client) ListSessions(ctx context.Context) ([]uuid.UUID, error) {
120-
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/sessions/list", nil)
121-
122-
resp, err := c.httpClient.Do(req)
118+
resp, err := amasshttp.RequestWebPage(ctx, c.httpClient, &amasshttp.Request{URL: c.base + "/sessions/list"})
123119
if err != nil {
124120
return nil, err
125121
}
126-
defer func() { _ = resp.Body.Close() }()
127122

128123
if resp.StatusCode != http.StatusOK {
129-
msg, err := readJSONError(resp)
124+
msg, err := readJSONError(resp.Body)
130125
if err != nil {
131126
return nil, fmt.Errorf("listSessions: status=%s", resp.Status)
132127
}
133128
return nil, fmt.Errorf("listSessions: status=%s error=%s", resp.Status, msg)
134129
}
135130

136131
var out ListSessionsResponse
137-
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
132+
if err := json.Unmarshal([]byte(resp.Body), &out); err != nil {
138133
return nil, err
139134
}
140135

@@ -151,16 +146,16 @@ func (c *Client) ListSessions(ctx context.Context) ([]uuid.UUID, error) {
151146

152147
// Terminates the session associated with the provided token.
153148
func (c *Client) TerminateSession(ctx context.Context, token uuid.UUID) error {
154-
req, _ := http.NewRequestWithContext(ctx, http.MethodDelete, c.base+"/sessions/"+token.String(), nil)
155-
156-
resp, err := c.httpClient.Do(req)
149+
resp, err := amasshttp.RequestWebPage(ctx, c.httpClient, &amasshttp.Request{
150+
Method: http.MethodDelete,
151+
URL: c.base + "/sessions/" + token.String(),
152+
})
157153
if err != nil {
158154
return err
159155
}
160-
defer func() { _ = resp.Body.Close() }()
161156

162157
if resp.StatusCode != http.StatusNoContent {
163-
msg, err := readJSONError(resp)
158+
msg, err := readJSONError(resp.Body)
164159
if err != nil {
165160
return fmt.Errorf("terminateSession: status=%s", resp.Status)
166161
}
@@ -171,24 +166,22 @@ func (c *Client) TerminateSession(ctx context.Context, token uuid.UUID) error {
171166

172167
// Retrieves statistics for the session associated with the provided token.
173168
func (c *Client) SessionStats(ctx context.Context, token uuid.UUID) (*et.SessionStats, error) {
174-
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/sessions/"+token.String()+"/stats", nil)
175-
176-
resp, err := c.httpClient.Do(req)
169+
resp, err := amasshttp.RequestWebPage(ctx, c.httpClient,
170+
&amasshttp.Request{URL: c.base + "/sessions/" + token.String() + "/stats"})
177171
if err != nil {
178172
return nil, err
179173
}
180-
defer func() { _ = resp.Body.Close() }()
181174

182175
if resp.StatusCode != http.StatusOK {
183-
msg, err := readJSONError(resp)
176+
msg, err := readJSONError(resp.Body)
184177
if err != nil {
185178
return nil, fmt.Errorf("%s/stats: status=%s", token.String(), resp.Status)
186179
}
187180
return nil, fmt.Errorf("%s/stats: status=%s error=%s", token.String(), resp.Status, msg)
188181
}
189182

190183
var st et.SessionStats
191-
if err := json.NewDecoder(resp.Body).Decode(&st); err != nil {
184+
if err := json.Unmarshal([]byte(resp.Body), &st); err != nil {
192185
return nil, err
193186
}
194187
return &st, nil
@@ -199,23 +192,24 @@ func (c *Client) SessionScope(ctx context.Context, token uuid.UUID, atype oam.As
199192
sessionID := token.String()
200193
atypestr := strings.ToLower(string(atype))
201194
u := fmt.Sprintf("%s/sessions/%s/scope/%s", c.base, sessionID, atypestr)
202-
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
203-
204-
resp, err := c.httpClient.Do(req)
195+
resp, err := amasshttp.RequestWebPage(ctx, c.httpClient, &amasshttp.Request{URL: u})
205196
if err != nil {
206197
return nil, err
207198
}
208-
defer func() { _ = resp.Body.Close() }()
209199

210200
if resp.StatusCode != http.StatusOK {
211-
msg, err := readJSONError(resp)
201+
msg, err := readJSONError(resp.Body)
212202
if err != nil {
213203
return nil, fmt.Errorf("%s/scope: status=%s", token.String(), resp.Status)
214204
}
215205
return nil, fmt.Errorf("%s/scope: status=%s error=%s", token.String(), resp.Status, msg)
216206
}
217207

218-
return apiclient.DecodeAssetsForScopeEndpoint(atype, resp.Body)
208+
reader := strings.NewReader(resp.Body)
209+
readCloser := io.NopCloser(reader)
210+
defer func() { _ = readCloser.Close() }()
211+
212+
return apiclient.DecodeAssetsForScopeEndpoint(atype, readCloser)
219213
}
220214

221215
// Creates a new asset on the server associated with the provided token.
@@ -228,25 +222,26 @@ func (c *Client) CreateAsset(ctx context.Context, token uuid.UUID, asset oam.Ass
228222

229223
sessionID := token.String()
230224
u := fmt.Sprintf("%s/sessions/%s/assets/%s", c.base, sessionID, atype)
231-
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, u, bytes.NewReader(raw))
232-
req.Header.Set("Content-Type", "application/json")
233-
234-
resp, err := c.httpClient.Do(req)
225+
resp, err := amasshttp.RequestWebPage(ctx, c.httpClient, &amasshttp.Request{
226+
URL: u,
227+
Body: string(raw),
228+
Method: http.MethodPost,
229+
Header: amasshttp.Header{"Content-Type": []string{"application/json"}},
230+
})
235231
if err != nil {
236232
return "", err
237233
}
238-
defer func() { _ = resp.Body.Close() }()
239234

240235
if resp.StatusCode != http.StatusOK {
241-
msg, err := readJSONError(resp)
236+
msg, err := readJSONError(resp.Body)
242237
if err != nil {
243238
return "", fmt.Errorf("createAsset: status=%s", resp.Status)
244239
}
245240
return "", fmt.Errorf("createAsset: status=%s error=%s", resp.Status, msg)
246241
}
247242

248243
var r AddAssetResponse
249-
if err := json.NewDecoder(resp.Body).Decode(&r); err != nil {
244+
if err := json.Unmarshal([]byte(resp.Body), &r); err != nil {
250245
return "", err
251246
}
252247
return r.EntityID, nil
@@ -277,28 +272,28 @@ func (c *Client) CreateAssetsBulk(ctx context.Context, token uuid.UUID, atype st
277272
}
278273

279274
sessionID := token.String()
280-
u := fmt.Sprintf("%s/sessions/%s/assets/%s:bulk", c.base, sessionID, atype)
281-
282275
body, _ := json.Marshal(BulkAddAssetsRequest{Items: items})
283-
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, u, bytes.NewReader(body))
284-
req.Header.Set("Content-Type", "application/json")
285-
286-
resp, err := c.httpClient.Do(req)
276+
u := fmt.Sprintf("%s/sessions/%s/assets/%s:bulk", c.base, sessionID, atype)
277+
resp, err := amasshttp.RequestWebPage(ctx, c.httpClient, &amasshttp.Request{
278+
URL: u,
279+
Body: string(body),
280+
Method: http.MethodPost,
281+
Header: amasshttp.Header{"Content-Type": []string{"application/json"}},
282+
})
287283
if err != nil {
288284
return 0, err
289285
}
290-
defer func() { _ = resp.Body.Close() }()
291286

292287
if resp.StatusCode != http.StatusOK {
293-
msg, err := readJSONError(resp)
288+
msg, err := readJSONError(resp.Body)
294289
if err != nil {
295290
return 0, fmt.Errorf("addAssetsBulk: status=%s", resp.Status)
296291
}
297292
return 0, fmt.Errorf("addAssetsBulk: status=%s error=%s", resp.Status, msg)
298293
}
299294

300295
var out BulkAddAssetsResponse
301-
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
296+
if err := json.Unmarshal([]byte(resp.Body), &out); err != nil {
302297
return 0, err
303298
}
304299
return int(out.Stored), nil
@@ -355,11 +350,11 @@ func (c *Client) Subscribe(ctx context.Context, token uuid.UUID) (<-chan string,
355350
return ch, nil
356351
}
357352

358-
func readJSONError(resp *http.Response) (string, error) {
353+
func readJSONError(content string) (string, error) {
359354
var errResp struct {
360355
Message string `json:"error"`
361356
}
362-
if err := json.NewDecoder(resp.Body).Decode(&errResp); err != nil {
357+
if err := json.Unmarshal([]byte(content), &errResp); err != nil {
363358
return "", err
364359
}
365360
return errResp.Message, nil

‎internal/enum/cli.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,7 @@ func CLIWorkflow(cmdName string, clArgs []string) {
233233
provFQDNs = append(provFQDNs, oamdns.FQDN{Name: a})
234234

235235
if fcount == client.MaxBulkItems {
236-
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
236+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
237237
defer cancel()
238238

239239
stored, err := c.CreateAssetsBulk(ctx, token, string(oam.FQDN), provFQDNs)
@@ -247,7 +247,7 @@ func CLIWorkflow(cmdName string, clArgs []string) {
247247
}
248248
}
249249
if fcount > 0 {
250-
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
250+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
251251
defer cancel()
252252

253253
if stored, err := c.CreateAssetsBulk(ctx, token, string(oam.FQDN), provFQDNs); err != nil {

‎internal/net/http/http.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -176,8 +176,8 @@ func RequestWebPage(ctx context.Context, client *http.Client, r *Request) (*Resp
176176
}
177177

178178
if r.Method == "" {
179-
r.Method = "GET"
180-
} else if r.Method != "GET" && r.Method != "POST" {
179+
r.Method = http.MethodGet
180+
} else if r.Method != http.MethodGet && r.Method != http.MethodPost && r.Method != http.MethodDelete {
181181
return nil, errors.New("failed to provide a valid HTTP method")
182182
}
183183

0 commit comments

Comments
 (0)