From 4470d6c18e798fb264e890d723b1a283a2f99752 Mon Sep 17 00:00:00 2001
From: Yongkun Gui <ygui@google.com>
Date: Mon, 21 Jun 2021 23:13:48 -0700
Subject: [PATCH] Add handler for /size/<number>

This handler will send an output of <number> bytes (all `x`s) back to
the client. This can be used to test return path of fragmented packets.

Signed-off-by: Yongkun Gui <ygui@google.com>
---
 main.go | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/main.go b/main.go
index 8d54475..a2d35eb 100644
--- a/main.go
+++ b/main.go
@@ -19,6 +19,8 @@ import (
 	"log"
 	"net"
 	"os"
+	"strconv"
+	"strings"
 	"text/template"
 
 	"pack.ag/tftp"
@@ -48,6 +50,21 @@ var tmpl = template.Must(template.New("response").Parse(response))
 
 func echoHandler(w tftp.ReadRequest) {
 	log.Printf("Read %q requested from %s\n", w.Name(), w.Addr())
+
+	// If name is /size/<number>, generate <number> * x as output
+	const prefix string = "size/"
+	if strings.HasPrefix(w.Name(), prefix) {
+		size, err := strconv.Atoi(strings.TrimPrefix(w.Name(), prefix))
+		if err != nil {
+			log.Printf("Malformatted size (should be a /size/<number>): %q\n", w.Name())
+			return
+		}
+
+		w.WriteSize(int64(size))
+		w.Write([]byte(strings.Repeat("x", size)))
+		return
+	}
+
 	err := tmpl.Execute(w, values{
 		Hostname:   hostname,
 		ClientAddr: w.Addr().IP,