fix: Use presigned URLs correctly

We need to use PutObject to actually upload the content to bucket. After that, we can use the presigned client to generate a URL that's accessible for 5 days. We also set the Content-Type correctly so browsers will correctly render the JSONL and HTML files that we upload. Signed-off-by: Dave Tucker <[email protected]>
instructlab · Mar 17, 2024 · 0cad5d9 · 0cad5d9
1 parent 4138b8c
commit 0cad5d9
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 9 deletions.
diff --git a/bot/src/worker.ts b/bot/src/worker.ts
@@ -42,7 +42,8 @@ export class Worker {
           issue_number: parseInt(prNumber),
           body:
             `Beep, boop 🤖  The test data has been generated!\n\n` +
-            `Find your results [here](${s3Url}).`,
+            `Find your results [here](${s3Url}).\n\n` +
+            `*This URL expires in 5 days.*`,
         };
 
         const octokit = await this.app.auth(parseInt(installationId));

diff --git a/worker/cmd/generate.go b/worker/cmd/generate.go
@@ -8,6 +8,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"path"
+	"strings"
 	"sync"
 	"syscall"
 	"text/template"
@@ -222,9 +223,10 @@ func processJob(ctx context.Context, conn redis.Conn, svc *s3.Client, logger *za
 		return
 	}
 
-	outDirName := fmt.Sprintf("generate-pr-%s-%s", prNumber, head.Name().Short())
+	outDirName := fmt.Sprintf("generate-pr-%s-%s", prNumber, head.Hash())
 	outputDir := path.Join(workDir, outDirName)
 
+	sugar = sugar.With("out_dir", outputDir)
 	_ = os.MkdirAll(outputDir, 0755)
 
 	lab := "lab"
@@ -254,26 +256,39 @@ func processJob(ctx context.Context, conn redis.Conn, svc *s3.Client, logger *za
 		return
 	}
 	presign := s3.NewPresignClient(svc, func(options *s3.PresignOptions) {
-		options.Expires = 24 * time.Hour * 30
+		// Must be less than a week
+		options.Expires = 24 * time.Hour * 5
 	})
 	presignedFiles := make([]map[string]string, 0)
 
 	for _, item := range items {
+		if strings.HasSuffix(item.Name(), "index.html") {
+			continue
+		}
 		file, err := os.Open(path.Join(outputDir, item.Name()))
 		if err != nil {
 			sugar.Errorf("Could not open file: %v", err)
 			return
 		}
 		defer file.Close()
+		_, err = svc.PutObject(ctx, &s3.PutObjectInput{
+			Bucket:      aws.String("instruct-lab-bot"),
+			Key:         aws.String(fmt.Sprintf("%s/%s", outDirName, item.Name())),
+			Body:        file,
+			ContentType: aws.String("application/json-lines+json"),
+		})
+		if err != nil {
+			sugar.Errorf("Could not upload file to S3: %v", err)
+			return
+		}
 
-		result, err := presign.PresignPutObject(ctx, &s3.PutObjectInput{
+		result, err := presign.PresignGetObject(ctx, &s3.GetObjectInput{
 			Bucket: aws.String("instruct-lab-bot"),
 			Key:    aws.String(fmt.Sprintf("%s/%s", outDirName, item.Name())),
-			Body:   file,
 		})
 
 		if err != nil {
-			sugar.Errorf("Could not presign object: %v", err)
+			sugar.Errorf("Could not generate presign get URL for object: %v", err)
 			return
 		}
 
@@ -288,17 +303,34 @@ func processJob(ctx context.Context, conn redis.Conn, svc *s3.Client, logger *za
 		sugar.Errorf("Could not create index.html: %v", err)
 		return
 	}
-	defer indexFile.Close()
 
 	if err := generateIndexHTML(indexFile, prNumber, presignedFiles); err != nil {
 		sugar.Errorf("Could not generate index.html: %v", err)
+		indexFile.Close()
+		return
+	}
+
+	indexFile, err = os.Open(indexFile.Name())
+	if err != nil {
+		sugar.Errorf("Could not re-read index.html: %v", err)
+		indexFile.Close()
+		return
+	}
+
+	_, err = svc.PutObject(ctx, &s3.PutObjectInput{
+		Bucket:      aws.String("instruct-lab-bot"),
+		Key:         aws.String(fmt.Sprintf("%s/index.html", outDirName)),
+		Body:        indexFile,
+		ContentType: aws.String("text/html"),
+	})
+	if err != nil {
+		sugar.Errorf("Could not upload index.html to S3: %v", err)
 		return
 	}
 
-	indexResult, err := presign.PresignPutObject(ctx, &s3.PutObjectInput{
+	indexResult, err := presign.PresignGetObject(ctx, &s3.GetObjectInput{
 		Bucket: aws.String("instruct-lab-bot"),
 		Key:    aws.String(fmt.Sprintf("%s/index.html", outDirName)),
-		Body:   indexFile,
 	})
 
 	if err != nil {