Merge pull request #5 from NextronSystems/feat/more-retries

add more retries to the different collector variants
NextronSystems · Jun 11, 2024 · 4c542aa · 4c542aa
2 parents bb5ee30 + e8aeb43
commit 4c542aa
Show file tree

Hide file tree

Showing 5 changed files with 121 additions and 66 deletions.
diff --git a/go/collector.go b/go/collector.go
@@ -239,7 +239,7 @@ func (c *Collector) uploadToThunderstorm(info *infoWithPath) (redo bool) {
 		if info.retries < 3 {
 			c.logger.Printf("Could not send file %s to thunderstorm, will try again: %v", info.path, err)
 			info.retries++
-			time.Sleep(time.Second)
+			time.Sleep(4 * time.Second * time.Duration(1<<info.retries))
 			return true
 		} else {
 			c.logger.Printf("Could not send file %s to thunderstorm, canceling it.", info.path)

diff --git a/scripts/thunderstorm-collector.pl b/scripts/thunderstorm-collector.pl
@@ -127,19 +127,32 @@ sub processDir {
 sub submitSample {
     my ($filepath) = shift;
     print "[SUBMIT] Submitting $filepath ...\n";
-    eval { 
-        my $req = $ua->post($api_endpoint,
-            Content_Type => 'form-data',
-            Content => [
-                "file" => [ $filepath ],
-            ],
-        );
-        $num_submitted++;
-        print "\nError: ", $req->status_line unless $req->is_success;
-    } or do {
-        my $error = $@ || 'Unknown failure';
-        warn "Could not submit '$filepath' - $error";
-    };
+    my $retry = 0;
+    for ($retry = 0; $retry < 4; $retry++) {
+        if ($retry > 0) {
+            my $sleep_time = 2 << $retry;
+            print "[SUBMIT] Waiting $sleep_time seconds to retry submitting $filepath ...\n";
+            sleep($sleep_time)
+        }
+        my $successful = 0;
+        eval {
+            my $req = $ua->post($api_endpoint,
+                Content_Type => 'form-data',
+                Content => [
+                    "file" => [ $filepath ],
+                ],
+            );
+            $successful = $req->is_success;
+            $num_submitted++;
+            print "\nError: ", $req->status_line unless $successful;
+        } or do {
+            my $error = $@ || 'Unknown failure';
+            warn "Could not submit '$filepath' - $error";
+        };
+        if ($successful) {
+            last;
+        }
+    }
 }
 
 # MAIN ----------------------------------------------------------------

diff --git a/scripts/thunderstorm-collector.ps1 b/scripts/thunderstorm-collector.ps1
@@ -274,6 +274,7 @@ try {
 
         # Submitting the request
         $StatusCode = 0
+        $Retries = 0
         while ( $($StatusCode) -ne 200 ) {
             try {
                 Write-Log "Submitting to Thunderstorm server: $($_.FullName) ..." -Level "Info"
@@ -291,8 +292,14 @@ try {
                     Write-Log "503: Server seems busy - retrying in $($WaitSecs) seconds"
                     Start-Sleep -Seconds $($WaitSecs)
                 } else {
-                    Write-Log "$($StatusCode): Server has problems - retrying in 3 seconds"
-                    Start-Sleep -Seconds 3
+                    if ( $Retries -eq 3) {
+                        Write-Log "$($StatusCode): Server still has problems - giving up"
+                        break
+                    }
+                    $Retries = $Retries + 1
+                    $SleepTime = 2 * [Math]::Pow(2, $Retries)
+                    Write-Log "$($StatusCode): Server has problems - retrying in $SleepTime seconds"
+                    Start-Sleep -Seconds $($SleepTime)
                 }
             }
         }

diff --git a/scripts/thunderstorm-collector.py b/scripts/thunderstorm-collector.py
@@ -5,9 +5,7 @@
 import os
 import re
 import ssl
-import sys
 import time
-import urllib.parse
 import uuid
 
 # Configuration
@@ -34,6 +32,11 @@
 num_submitted = 0
 num_processed = 0
 
+# URL to use for submission
+api_endpoint = ""
+
+# Original args
+args = {}
 
 # Functions
 def process_dir(workdir):
@@ -106,53 +109,70 @@ def skip_file(filepath):
 def submit_sample(filepath):
     print("[SUBMIT] Submitting {} ...".format(filepath))
 
+    headers = {
+        "Content-Type": "application/octet-stream",
+        "Content-Disposition": f"attachment; filename={filepath}",
+    }
+
     try:
-        headers = {
-            "Content-Type": "application/octet-stream",
-            "Content-Disposition": f"attachment; filename={filepath}",
-        }
 
         with open(filepath, "rb") as f:
             data = f.read()
 
-        boundary = str(uuid.uuid4())
-        headers = {
-            "Content-Type": f"multipart/form-data; boundary={boundary}",
-        }
-
-        # Create multipart/form-data payload
-        payload = (
-            f"--{boundary}\r\n"
-            f'Content-Disposition: form-data; name="file"; filename="{filepath}"\r\n'
-            f"Content-Type: application/octet-stream\r\n\r\n"
-        ).encode("utf-8")
-        payload += data
-        payload += f"\r\n--{boundary}--\r\n".encode("utf-8")
-
-        if args.tls:
-            if args.insecure:
-                context = ssl._create_unverified_context()
+    except Exception as e:
+        print("[ERROR] Could not read '{}' - {}".format(filepath, e))
+        return
+
+    boundary = str(uuid.uuid4())
+    headers = {
+        "Content-Type": f"multipart/form-data; boundary={boundary}",
+    }
+
+    # Create multipart/form-data payload
+    payload = (
+        f"--{boundary}\r\n"
+        f'Content-Disposition: form-data; name="file"; filename="{filepath}"\r\n'
+        f"Content-Type: application/octet-stream\r\n\r\n"
+    ).encode("utf-8")
+    payload += data
+    payload += f"\r\n--{boundary}--\r\n".encode("utf-8")
+
+    retries = 0
+    while retries < 3:
+        try:
+            if args.tls:
+                if args.insecure:
+                    context = ssl._create_unverified_context()
+                else:
+                    context = ssl._create_default_https_context()
+                conn = http.client.HTTPSConnection(args.server, args.port, context=context)
             else:
-                context = ssl._create_default_https_context()
-            conn = http.client.HTTPSConnection(args.server, args.port, context=context)
-        else:
-            conn = http.client.HTTPConnection(args.server, args.port)
-        conn.request("POST", api_endpoint, body=payload, headers=headers)
+                conn = http.client.HTTPConnection(args.server, args.port)
+            conn.request("POST", api_endpoint, body=payload, headers=headers)
 
-        resp = conn.getresponse()
+            resp = conn.getresponse()
 
-        global num_submitted
-        num_submitted += 1
+        except Exception as e:
+            print("[ERROR] Could not submit '{}' - {}".format(filepath, e))
+            retries += 1
+            time.sleep(2 << retries)
+            continue
 
-        if resp.status != 200:
-            print(
-                "[ERROR] HTTP return status: {}, reason: {}".format(
-                    resp.status, resp.reason
-                )
+        # pylint: disable=no-else-continue
+        if resp.status == 503: # Service unavailable
+            retry_time = resp.headers.get("Retry-After", 30)
+            time.sleep(retry_time)
+            continue
+        elif resp.status == 200:
+            break
+        print(
+            "[ERROR] HTTP return status: {}, reason: {}".format(
+                resp.status, resp.reason
             )
+        )
 
-    except Exception as e:
-        print("Could not submit '{}' - {}".format(filepath, e))
+    global num_submitted
+    num_submitted += 1
 
 
 # Main
@@ -188,13 +208,11 @@ def submit_sample(filepath):
     )
     parser.add_argument("--debug", action="store_true", help="Enable debug logging.")
 
-    global args
     args = parser.parse_args()
 
     if args.tls:
         schema = "https"
 
-    global api_endpoint
     api_endpoint = "{}://{}:{}/api/checkAsync".format(schema, args.server, args.port)
 
     print("=" * 80)
@@ -214,8 +232,8 @@ def submit_sample(filepath):
     print("Starting the walk at: {} ...".format(", ".join(args.dirs)))
 
     # Walk directory
-    for dir in args.dirs:
-        process_dir(dir)
+    for walkdir in args.dirs:
+        process_dir(walkdir)
 
     # End message
     end_date = time.time()

diff --git a/scripts/thunderstorm-collector.sh b/scripts/thunderstorm-collector.sh
@@ -137,14 +137,31 @@ do
                 continue
             fi
             log debug "Submitting ${file_path} ..."
-            # Submit sample
-            result=$(curl -s -X POST \
-                     "$scheme://$THUNDERSTORM_SERVER:8080/api/$api_endpoint" \
-                     --form "file=@${file_path};filename=${file_path}")
-            # If not 'id' in result
-            error="reason"
-            if [ "${result/$error}" != "$result" ]; then
-                log error "$result"
+            successful=0
+
+            for retry in {1..3}; do
+                # Submit sample
+                result=$(curl -s -X POST \
+                        "$scheme://$THUNDERSTORM_SERVER:8080/api/$api_endpoint" \
+                        --form "file=@${file_path};filename=${file_path}")
+                curl_exit=$?
+                if [ $curl_exit -ne 0 ]; then
+                    log error "Upload failed with code $curl_exit"
+                    sleep $((2 << retry))
+                    continue
+                fi
+
+                # If 'reason' in result
+                if [ "${result/reason}" != "$result" ]; then
+                    log error "$result"
+                    sleep $((2 << retry))
+                    continue
+                fi
+                successful=1
+                break
+            done
+            if [ $successful -ne 1 ]; then
+                log error "Could not upload ${file_path}"
             fi
         fi
     done