From 621a23693a4dfa3b7f8d1147fd4b48d535403019 Mon Sep 17 00:00:00 2001
From: "Jeffrey N. Johnson" <jeff@cohere-llc.com>
Date: Tue, 23 Jul 2024 14:36:46 -0700
Subject: [PATCH] Added a maximum payload size constraint.

---
 config/config.go |  4 ++++
 dts.yaml.example |  1 +
 tasks/tasks.go   | 25 +++++++++++++++++++++++++
 3 files changed, 30 insertions(+)

diff --git a/config/config.go b/config/config.go
index f304780..ee262a8 100644
--- a/config/config.go
+++ b/config/config.go
@@ -37,6 +37,9 @@ type serviceConfig struct {
 	// maximum number of allowed incoming connections
 	// default: 100
 	MaxConnections int `json:"max_connections,omitempty" yaml:"max_connections,omitempty"`
+	// maximum size of requested payload for transfer, past which transfer
+	// requests are rejected (gigabytes)
+	MaxPayloadSize int `json:"max_payload_size,omitempy" yaml:"max_payload_size,omitempty"`
 	// polling interval for checking transfer statuses (milliseconds)
 	// default: 1 minute
 	PollInterval int `json:"poll_interval" yaml:"poll_interval"`
@@ -80,6 +83,7 @@ func readConfig(bytes []byte) error {
 	var conf configFile
 	conf.Service.Port = 8080
 	conf.Service.MaxConnections = 100
+	conf.Service.MaxPayloadSize = 100
 	conf.Service.PollInterval = int(time.Minute / time.Millisecond)
 	conf.Service.DeleteAfter = 7 * 24 * 3600
 	err := yaml.Unmarshal(bytes, &conf)
diff --git a/dts.yaml.example b/dts.yaml.example
index 61d79e9..a252528 100644
--- a/dts.yaml.example
+++ b/dts.yaml.example
@@ -5,6 +5,7 @@
 service:
   port: 8080                 # port on which the service listenѕ
   max_connections: 100       # maximum number of incoming HTTP connections
+  max_payload_size: 100      # limit (if any) on DTS payload size (gigabytes)
   poll_interval:   60000     # interval at which DTS checks transfer statuses (ms)
   endpoint: globus-local     # name of endpoint used for manifest generation
   data_dir: /path/to/dir     # directory DTS uses for internal data storage
diff --git a/tasks/tasks.go b/tasks/tasks.go
index 1ebe9ce..9ff75fb 100644
--- a/tasks/tasks.go
+++ b/tasks/tasks.go
@@ -81,6 +81,25 @@ type taskType struct {
 	CompletionTime      time.Time      // time at which the transfer completed
 }
 
+// This error type is returned when a payload is requested that is too large.
+type PayloadTooLargeError struct {
+	size int // size of the requested payload in gigabytes
+}
+
+func (e PayloadTooLargeError) Error() string {
+	return fmt.Sprintf("Requested payload is too large: %d GB (limit is %d GB).",
+		e.size, config.Service.MaxPayloadSize)
+}
+
+// computes the size of a payload for a transfer task (in gigabytes)
+func payloadSize(resources []DataResource) int {
+	var size uint64
+	for _, resource := range resources {
+		size += uint64(resource.Bytes)
+	}
+	return int(size / (1024 * 1024))
+}
+
 // starts a task going, initiating staging if needed
 func (task *taskType) start() error {
 	source, err := databases.NewDatabase(task.Orcid, task.Source)
@@ -94,6 +113,12 @@ func (task *taskType) start() error {
 		return err
 	}
 
+	// make sure the size of the payload doesn't exceed our specified limit
+	size := payloadSize(task.Resources) // (in GB)
+	if size > config.Service.MaxPayloadSize {
+		return &PayloadTooLargeError{size: size}
+	}
+
 	// are the files already staged? (only works for public data)
 	sourceEndpoint, err := source.Endpoint()
 	if err != nil {