From a64d839e2c17c2dc91e5d68321896bfaee3ab838 Mon Sep 17 00:00:00 2001
From: Yuhui Xu <xuyh0120@outlook.com>
Date: Thu, 2 Sep 2021 20:21:28 -0500
Subject: [PATCH] frontend: limit fetched response size to 64KB (#39)

---
 frontend/lgproxy.go | 12 +++++++++---
 frontend/whois.go   | 10 ++++++----
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/frontend/lgproxy.go b/frontend/lgproxy.go
index a79aed5..99c9620 100644
--- a/frontend/lgproxy.go
+++ b/frontend/lgproxy.go
@@ -1,7 +1,7 @@
 package main
 
 import (
-	"io/ioutil"
+	"io"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -53,8 +53,14 @@ func batchRequest(servers []string, endpoint string, command string) []string {
 					ch <- channelData{i, "request failed: " + err.Error() + "\n"}
 					return
 				}
-				text, _ := ioutil.ReadAll(response.Body)
-				ch <- channelData{i, string(text)}
+
+				buf := make([]byte, 65536)
+				_, err = io.ReadFull(response.Body, buf)
+				if err != nil && err != io.ErrUnexpectedEOF {
+					ch <- channelData{i, err.Error()}
+				} else {
+					ch <- channelData{i, string(buf)}
+				}
 			}(url, i)
 		}
 	}
diff --git a/frontend/whois.go b/frontend/whois.go
index b4cc173..cd17d47 100644
--- a/frontend/whois.go
+++ b/frontend/whois.go
@@ -1,7 +1,7 @@
 package main
 
 import (
-	"io/ioutil"
+	"io"
 	"net"
 	"time"
 )
@@ -19,9 +19,11 @@ func whois(s string) string {
 	defer conn.Close()
 
 	conn.Write([]byte(s + "\r\n"))
-	result, err := ioutil.ReadAll(conn)
-	if err != nil {
+
+	buf := make([]byte, 65536)
+	_, err = io.ReadFull(conn, buf)
+	if err != nil && err != io.ErrUnexpectedEOF {
 		return err.Error()
 	}
-	return string(result)
+	return string(buf)
 }