Skip to content

Commit b066586

Browse files
secDre4mersecDre4mer
committed
feat: add retry logic to more collectors
1 parent b2dc44f commit b066586

4 files changed

+112
-59
lines changed

Diff for: ‎scripts/thunderstorm-collector.pl

100644100755
+26-13
Original file line numberDiff line numberDiff line change
@@ -127,19 +127,32 @@ sub processDir {
127127
sub submitSample {
128128
my ($filepath) = shift;
129129
print "[SUBMIT] Submitting $filepath ...\n";
130-
eval {
131-
my $req = $ua->post($api_endpoint,
132-
Content_Type => 'form-data',
133-
Content => [
134-
"file" => [ $filepath ],
135-
],
136-
);
137-
$num_submitted++;
138-
print "\nError: ", $req->status_line unless $req->is_success;
139-
} or do {
140-
my $error = $@ || 'Unknown failure';
141-
warn "Could not submit '$filepath' - $error";
142-
};
130+
my $retry = 0;
131+
for ($retry = 0; $retry < 4; $retry++) {
132+
if ($retry > 0) {
133+
my $sleep_time = 2 << $retry;
134+
print "[SUBMIT] Waiting $sleep_time seconds to retry submitting $filepath ...\n";
135+
sleep($sleep_time)
136+
}
137+
my $successful = 0;
138+
eval {
139+
my $req = $ua->post($api_endpoint,
140+
Content_Type => 'form-data',
141+
Content => [
142+
"file" => [ $filepath ],
143+
],
144+
);
145+
$successful = $req->is_success;
146+
$num_submitted++;
147+
print "\nError: ", $req->status_line unless $successful;
148+
} or do {
149+
my $error = $@ || 'Unknown failure';
150+
warn "Could not submit '$filepath' - $error";
151+
};
152+
if ($successful) {
153+
last;
154+
}
155+
}
143156
}
144157

145158
# MAIN ----------------------------------------------------------------

Diff for: ‎scripts/thunderstorm-collector.ps1

+9-2
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,7 @@ try {
274274

275275
# Submitting the request
276276
$StatusCode = 0
277+
$Retries = 0
277278
while ( $($StatusCode) -ne 200 ) {
278279
try {
279280
Write-Log "Submitting to Thunderstorm server: $($_.FullName) ..." -Level "Info"
@@ -291,8 +292,14 @@ try {
291292
Write-Log "503: Server seems busy - retrying in $($WaitSecs) seconds"
292293
Start-Sleep -Seconds $($WaitSecs)
293294
} else {
294-
Write-Log "$($StatusCode): Server has problems - retrying in 3 seconds"
295-
Start-Sleep -Seconds 3
295+
if ( $Retries -eq 3) {
296+
Write-Log "$($StatusCode): Server still has problems - giving up"
297+
break
298+
}
299+
$Retries = $Retries + 1
300+
$SleepTime = 2 * [Math]::Pow(2, $Retries)
301+
Write-Log "$($StatusCode): Server has problems - retrying in $SleepTime seconds"
302+
Start-Sleep -Seconds $($SleepTime)
296303
}
297304
}
298305
}

Diff for: ‎scripts/thunderstorm-collector.py

100644100755
+52-36
Original file line numberDiff line numberDiff line change
@@ -106,53 +106,69 @@ def skip_file(filepath):
106106
def submit_sample(filepath):
107107
print("[SUBMIT] Submitting {} ...".format(filepath))
108108

109+
headers = {
110+
"Content-Type": "application/octet-stream",
111+
"Content-Disposition": f"attachment; filename={filepath}",
112+
}
113+
109114
try:
110-
headers = {
111-
"Content-Type": "application/octet-stream",
112-
"Content-Disposition": f"attachment; filename={filepath}",
113-
}
114115

115116
with open(filepath, "rb") as f:
116117
data = f.read()
117118

118-
boundary = str(uuid.uuid4())
119-
headers = {
120-
"Content-Type": f"multipart/form-data; boundary={boundary}",
121-
}
122-
123-
# Create multipart/form-data payload
124-
payload = (
125-
f"--{boundary}\r\n"
126-
f'Content-Disposition: form-data; name="file"; filename="{filepath}"\r\n'
127-
f"Content-Type: application/octet-stream\r\n\r\n"
128-
).encode("utf-8")
129-
payload += data
130-
payload += f"\r\n--{boundary}--\r\n".encode("utf-8")
131-
132-
if args.tls:
133-
if args.insecure:
134-
context = ssl._create_unverified_context()
119+
except Exception as e:
120+
print("[ERROR] Could not read '{}' - {}".format(filepath, e))
121+
return
122+
123+
boundary = str(uuid.uuid4())
124+
headers = {
125+
"Content-Type": f"multipart/form-data; boundary={boundary}",
126+
}
127+
128+
# Create multipart/form-data payload
129+
payload = (
130+
f"--{boundary}\r\n"
131+
f'Content-Disposition: form-data; name="file"; filename="{filepath}"\r\n'
132+
f"Content-Type: application/octet-stream\r\n\r\n"
133+
).encode("utf-8")
134+
payload += data
135+
payload += f"\r\n--{boundary}--\r\n".encode("utf-8")
136+
137+
retries = 0
138+
while retries < 3:
139+
try:
140+
if args.tls:
141+
if args.insecure:
142+
context = ssl._create_unverified_context()
143+
else:
144+
context = ssl._create_default_https_context()
145+
conn = http.client.HTTPSConnection(args.server, args.port, context=context)
135146
else:
136-
context = ssl._create_default_https_context()
137-
conn = http.client.HTTPSConnection(args.server, args.port, context=context)
138-
else:
139-
conn = http.client.HTTPConnection(args.server, args.port)
140-
conn.request("POST", api_endpoint, body=payload, headers=headers)
147+
conn = http.client.HTTPConnection(args.server, args.port)
148+
conn.request("POST", api_endpoint, body=payload, headers=headers)
141149

142-
resp = conn.getresponse()
150+
resp = conn.getresponse()
143151

144-
global num_submitted
145-
num_submitted += 1
152+
except Exception as e:
153+
print("[ERROR] Could not submit '{}' - {}".format(filepath, e))
154+
retries += 1
155+
time.sleep(2 << retries)
156+
continue
146157

147-
if resp.status != 200:
148-
print(
149-
"[ERROR] HTTP return status: {}, reason: {}".format(
150-
resp.status, resp.reason
151-
)
158+
if resp.status == 503: # Service unavailable
159+
retry_time = resp.headers.get("Retry-After", 30)
160+
time.sleep(retry_time)
161+
continue
162+
elif resp.status == 200:
163+
break
164+
print(
165+
"[ERROR] HTTP return status: {}, reason: {}".format(
166+
resp.status, resp.reason
152167
)
168+
)
153169

154-
except Exception as e:
155-
print("Could not submit '{}' - {}".format(filepath, e))
170+
global num_submitted
171+
num_submitted += 1
156172

157173

158174
# Main

Diff for: ‎scripts/thunderstorm-collector.sh

+25-8
Original file line numberDiff line numberDiff line change
@@ -137,14 +137,31 @@ do
137137
continue
138138
fi
139139
log debug "Submitting ${file_path} ..."
140-
# Submit sample
141-
result=$(curl -s -X POST \
142-
"$scheme://$THUNDERSTORM_SERVER:8080/api/$api_endpoint" \
143-
--form "file=@${file_path};filename=${file_path}")
144-
# If not 'id' in result
145-
error="reason"
146-
if [ "${result/$error}" != "$result" ]; then
147-
log error "$result"
140+
successful=0
141+
142+
for retry in {1..3}; do
143+
# Submit sample
144+
result=$(curl -s -X POST \
145+
"$scheme://$THUNDERSTORM_SERVER:8080/api/$api_endpoint" \
146+
--form "file=@${file_path};filename=${file_path}")
147+
curl_exit=$?
148+
if [ $curl_exit -ne 0 ]; then
149+
log error "Upload failed with code $curl_exit"
150+
sleep $((2 << retry))
151+
continue
152+
fi
153+
154+
# If 'reason' in result
155+
if [ "${result/reason}" != "$result" ]; then
156+
log error "$result"
157+
sleep $((2 << retry))
158+
continue
159+
fi
160+
successful=1
161+
break
162+
done
163+
if [ $successful -ne 1 ]; then
164+
log error "Could not upload ${file_path}"
148165
fi
149166
fi
150167
done

0 commit comments

Comments
 (0)