Reuse presigned URLS (#977)

* reused Presignedurls * add time checking * shorten minutes to test * add logs * convert date string to date * fix comparision * fix logic * restructure logic * remove logs * change minutes back to 30 * reuse dup code
ashirt-ops · Nov 21, 2023 · 8fd6b29 · 8fd6b29
1 parent 798584c
commit 8fd6b29
Show file tree

Hide file tree

Showing 7 changed files with 80 additions and 33 deletions.
diff --git a/backend/contentstore/s3store.go b/backend/contentstore/s3store.go
@@ -71,20 +71,30 @@ func (s *S3Store) Read(key string) (io.Reader, error) {
 	return res.Body, nil
 }
 
-func (s *S3Store) SendURL(key string) (*string, error) {
+type URLData struct {
+	Url            string    `json:"url"`
+	ExpirationTime time.Time `json:"expirationTime"`
+}
+
+func (s *S3Store) SendURLData(key string) (*URLData, error) {
 	contentType := "image/jpeg"
 	req, _ := s.s3Client.GetObjectRequest(&s3.GetObjectInput{
 		Bucket:              aws.String(s.bucketName),
 		Key:                 aws.String(key),
 		ResponseContentType: aws.String(contentType),
 	})
 
-	url, err := req.Presign(time.Minute * 30)
+	minutes := time.Minute * time.Duration(30)
+	url, err := req.Presign(minutes)
 	if err != nil {
 		return nil, backend.WrapError("Unable to get presigned URL", err)
 	}
+	data := URLData{
+		Url:            url,
+		ExpirationTime: time.Now().UTC().Add(minutes),
+	}
 
-	return &url, nil
+	return &data, nil
 }
 
 // Delete removes files in in your OS's temp directory

diff --git a/backend/server/web.go b/backend/server/web.go
@@ -521,11 +521,15 @@ func bindWebRoutes(r chi.Router, db *database.Connection, contentStore contentst
 			return nil, backend.WrapError("Unable to read evidence", err)
 		}
 		if s3Store, ok := contentStore.(*contentstore.S3Store); ok && evidence.ContentType == "image" {
-			url, err := services.SendUrl(r.Context(), db, s3Store, i)
+			urlData, err := services.SendURLData(r.Context(), db, s3Store, i)
 			if err != nil {
-				return nil, backend.WrapError("Unable get s3 URL", err)
+				return nil, backend.WrapError("Unable to get s3 URL", err)
 			}
-			return bytes.NewReader([]byte(*url)), nil
+			jsonifiedData, err := json.Marshal(urlData)
+			if err != nil {
+				return nil, backend.WrapError("Unable to send s3 URL", err)
+			}
+			return bytes.NewReader(jsonifiedData), nil
 		}
 		if i.LoadPreview {
 			return evidence.Preview, nil

diff --git a/backend/services/evidence.go b/backend/services/evidence.go
@@ -329,9 +329,9 @@ func ListEvidenceForOperation(ctx context.Context, db *database.Connection, cont
 			tags = []dtos.Tag{}
 		}
 
-		sendUrl := false
+		sendURL := false
 		if usingS3 && evi.ContentType == "image" {
-			sendUrl = true
+			sendURL = true
 		}
 
 		evidenceDTO[idx] = &dtos.Evidence{
@@ -341,26 +341,26 @@ func ListEvidenceForOperation(ctx context.Context, db *database.Connection, cont
 			OccurredAt:  evi.OccurredAt,
 			ContentType: evi.ContentType,
 			Tags:        tags,
-			SendUrl:     sendUrl,
+			SendUrl:     sendURL,
 		}
 	}
 	return evidenceDTO, nil
 }
 
-func SendUrl(ctx context.Context, db *database.Connection, contentStore *contentstore.S3Store, i ReadEvidenceInput) (*string, error) {
+func SendURLData(ctx context.Context, db *database.Connection, contentStore *contentstore.S3Store, i ReadEvidenceInput) (*contentstore.URLData, error) {
 	operation, evidence, err := lookupOperationEvidence(db, i.OperationSlug, i.EvidenceUUID)
 	if err != nil {
 		return nil, backend.WrapError("Unable to read evidence", backend.UnauthorizedReadErr(err))
 	}
 	if err := policy.Require(middleware.Policy(ctx), policy.CanReadOperation{OperationID: operation.ID}); err != nil {
 		return nil, backend.WrapError("Unwilling to read evidence", backend.UnauthorizedReadErr(err))
 	}
-	str, err := contentStore.SendURL(evidence.FullImageKey)
+	urlData, err := contentStore.SendURLData(evidence.FullImageKey)
 	if err != nil {
 		return nil, backend.WrapError("Unable to get image URL", backend.ServerErr(err))
 	}
 
-	return str, nil
+	return urlData, nil
 
 }
 func ReadEvidence(ctx context.Context, db *database.Connection, contentStore contentstore.Store, i ReadEvidenceInput) (*ReadEvidenceOutput, error) {

diff --git a/frontend/src/components/evidence_preview/index.tsx b/frontend/src/components/evidence_preview/index.tsx
@@ -5,8 +5,8 @@ import * as React from 'react'
 import classnames from 'classnames/bind'
 import { CodeBlockViewer } from '../code_block'
 import { HarViewer, isAHar } from '../http_cycle_viewer'
-import { SupportedEvidenceType, CodeBlock, EvidenceViewHint, InteractionHint } from 'src/global_types'
-import { getEvidenceAsCodeblock, getEvidenceAsString, updateEvidence } from 'src/services/evidence'
+import { SupportedEvidenceType, CodeBlock, EvidenceViewHint, InteractionHint, UrlData } from 'src/global_types'
+import { getEvidenceAsCodeblock, getEvidenceAsString, getEvidenceAsUrlData, updateEvidence } from 'src/services/evidence'
 import { useWiredData } from 'src/helpers'
 import ErrorDisplay from 'src/components/error_display'
 import LazyLoadComponent from 'src/components/lazy_load_component'
@@ -44,6 +44,8 @@ export default (props: {
   fitToContainer?: boolean,
   useS3Url: boolean,
   onClick?: (event: React.MouseEvent<HTMLDivElement, MouseEvent>) => void,
+  urlSetter?: (urlData: UrlData | null) => void,
+  preSavedS3UrlData?: UrlData,
 }) => {
   const Component = getComponent(props.contentType)
   if (Component == null) return null
@@ -69,6 +71,8 @@ type EvidenceProps = {
   viewHint?: EvidenceViewHint,
   interactionHint?: InteractionHint,
   useS3Url: boolean
+  urlSetter?: (urlData: UrlData) => void,
+  preSavedS3UrlData?: UrlData,
 }
 
 const EvidenceCodeblock = (props: EvidenceProps) => {
@@ -81,16 +85,21 @@ const EvidenceCodeblock = (props: EvidenceProps) => {
 }
 
 const EvidenceImage = (props: EvidenceProps) => {
-  if (props.useS3Url) {
-    const wiredUrl = useWiredData<string>(React.useCallback(() => getEvidenceAsString({
+  let url = `/web/operations/${props.operationSlug}/evidence/${props.evidenceUuid}/media`
+  const now = new Date()
+  if (props.useS3Url && props.preSavedS3UrlData && new Date(props.preSavedS3UrlData.expirationTime) > now){
+    url = props.preSavedS3UrlData.url
+  } else if (props.useS3Url) {
+    const wiredUrl = useWiredData<UrlData>(React.useCallback(() => getEvidenceAsUrlData({
       operationSlug: props.operationSlug,
       evidenceUuid: props.evidenceUuid,
     }), [props.operationSlug, props.evidenceUuid]))
-    return wiredUrl.render(url => <img src={url} />)
-  } else {
-    const fullUrl = `/web/operations/${props.operationSlug}/evidence/${props.evidenceUuid}/media`
-    return <img src={fullUrl} />
+    wiredUrl.expose(s3url => {
+      props.urlSetter && props.urlSetter(s3url)
+      url = s3url.url
+    })
   }
+  return <img src={url} />
 }
 
 const EvidenceEvent = (_props: EvidenceProps) => {

diff --git a/frontend/src/components/timeline/index.tsx b/frontend/src/components/timeline/index.tsx
@@ -9,7 +9,7 @@ import TagList from 'src/components/tag_list'
 import classnames from 'classnames/bind'
 import Help from 'src/components/help'
 import { ClickPopover } from 'src/components/popover'
-import { Tag, Evidence } from 'src/global_types'
+import { Tag, Evidence, UrlData } from 'src/global_types'
 import { addTagToQuery, addOperatorToQuery } from 'src/helpers'
 import { default as Button, ButtonGroup } from 'src/components/button'
 import { CopyTextButton } from 'src/components/text_copiers'
@@ -39,6 +39,7 @@ export default (props: {
 
   const [activeChildIndex, setActiveChildIndex] = React.useState<number>(0)
   const [quicklookVisible, setQuicklookVisible] = React.useState<boolean>(false)
+  const [currImageUrlData, setCurrImageUrlData] = React.useState<UrlData| null>(null)
 
   const onKeyDown = (e: KeyboardEvent) => {
     // Only handle keystrokes if nothing is focused (target is body)
@@ -88,17 +89,21 @@ export default (props: {
 
   return <>
     <div className={cx('root')} ref={rootRef}>
-      {props.evidence.map((evi, idx) => (
-        <TimelineRow
-          {...props}
-          focusUuid={props.scrollToUuid}
-          active={activeChildIndex === idx}
-          evidence={evi}
-          key={evi.uuid}
-          onPreviewClick={() => { setActiveChildIndex(idx); setQuicklookVisible(true) }}
-          onClick={() => setActiveChildIndex(idx)}
-        />
-      ))}
+      {props.evidence.map((evi, idx) => {
+        const active = activeChildIndex === idx
+        return (
+          <TimelineRow
+            {...props}
+            focusUuid={props.scrollToUuid}
+            active={active}
+            urlSetter={active ? setCurrImageUrlData : undefined}
+            evidence={evi}
+            key={evi.uuid}
+            onPreviewClick={() => { setActiveChildIndex(idx); setQuicklookVisible(true) }}
+            onClick={() => setActiveChildIndex(idx)}
+          />
+        )
+      })}
       <Help className={cx('help')}
         preamble="Review and Edit the accumulated evidence for this operation"
         shortcuts={KeyboardShortcuts}
@@ -112,6 +117,7 @@ export default (props: {
           evidenceUuid={activeEvidence.uuid}
           contentType={activeEvidence.contentType}
           useS3Url={activeEvidence.sendUrl}
+          preSavedS3UrlData={currImageUrlData ? currImageUrlData : undefined}
           viewHint="large"
           interactionHint="active"
         />
@@ -131,6 +137,7 @@ const TimelineRow = (props: {
   focusUuid?: string,
   onPreviewClick: () => void,
   onClick: () => void,
+  urlSetter?: (urlData: UrlData | null) => void,
 }) => {
   const self = React.useRef<null | HTMLDivElement>(null)
 
@@ -160,6 +167,7 @@ const TimelineRow = (props: {
           evidenceUuid={props.evidence.uuid}
           contentType={props.evidence.contentType}
           useS3Url={props.evidence.sendUrl}
+          urlSetter={props.urlSetter}
           viewHint="medium"
           interactionHint="inactive"
         />

diff --git a/frontend/src/global_types.ts b/frontend/src/global_types.ts
@@ -356,3 +356,8 @@ export type OperationVariableData = {
   operationSlug: string,
   variable: OperationVar,
 }
+
+export type UrlData = {
+  url: string,
+  expirationTime: Date,
+}
diff --git a/frontend/src/services/evidence.ts b/frontend/src/services/evidence.ts
@@ -1,4 +1,4 @@
-import { Evidence, Finding, Tag, SubmittableEvidence, CodeBlock, TagDifference, EvidenceMetadata } from 'src/global_types'
+import { Evidence, Finding, Tag, SubmittableEvidence, CodeBlock, TagDifference, EvidenceMetadata, UrlData } from 'src/global_types'
 import { backendDataSource as ds } from './data_sources/backend'
 import { computeDelta } from 'src/helpers'
 import { evidenceFromDto } from './data_sources/converters'
@@ -40,6 +40,17 @@ export async function getEvidenceAsCodeblock(i: {
   }
 }
 
+export async function getEvidenceAsUrlData(i: {
+  operationSlug: string,
+  evidenceUuid: string,
+}): Promise<UrlData> {
+  const evi = JSON.parse(await ds.readEvidenceContent(i))
+  return {
+    url: evi.url,
+    expirationTime: evi.expirationTime,
+  }
+}
+
 export async function getEvidenceAsString(i: {
   operationSlug: string,
   evidenceUuid: string,