1 files changed, 20 insertions, 9 deletions
diff --git a/nslcd_server/func_handlerequest.go.gen b/nslcd_server/func_handlerequest.go.gen
index af36e84..7c28e7c 100755
--- a/nslcd_server/func_handlerequest.go.gen
+++ b/nslcd_server/func_handlerequest.go.gen
@@ -26,6 +26,7 @@ package nslcd_server
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"time"
 
@@ -53,6 +54,10 @@ type Limits struct {
 
 	// How long can we spend writing a response?
 	WriteTimeout time.Duration
+
+	// What is the maximum request length in bytes that we are
+	// willing to handle?
+	RequestMaxSize int64
 }
 
 type Conn interface {
@@ -92,13 +97,19 @@ func HandleRequest(backend Backend, limits Limits, conn Conn, cred unix.Ucred) (
 		}
 	}
 
+	var in io.Reader = conn
+	if limits.RequestMaxSize > 0 {
+		in = &io.LimitedReader{R: in, N: limits.RequestMaxSize}
+	}
+	out := conn
+
 	var version int32
-	maybePanic(p.Read(conn, &version))
+	maybePanic(p.Read(in, &version))
 	if version != p.NSLCD_VERSION {
 		return p.NslcdError(fmt.Sprintf("Version mismatch: server=%#08x client=%#08x", p.NSLCD_VERSION, version))
 	}
 	var action int32
-	maybePanic(p.Read(conn, &action))
+	maybePanic(p.Read(in, &action))
 
 	switch action {
 $(
@@ -106,7 +117,7 @@ while read -r request; do
 	cat <<EOT
 	case p.NSLCD_ACTION_${request^^}:
 		var req p.Request_${request}
-		maybePanic(p.Read(conn, &req))
+		maybePanic(p.Read(in, &req))
 		$(
 		case "$request" in
 			PAM_Authentication)
@@ -133,24 +144,24 @@ while read -r request; do
 		esac
 		)
 		if limits.WriteTimeout != 0 {
-			err = conn.SetWriteDeadline(time.Now().Add(limits.WriteTimeout))
+			err = out.SetWriteDeadline(time.Now().Add(limits.WriteTimeout))
 			if err != nil {
 				return err
 			}
 		}
-		maybePanic(p.Write(conn, p.NSLCD_VERSION))
-		maybePanic(p.Write(conn, action))
+		maybePanic(p.Write(out, p.NSLCD_VERSION))
+		maybePanic(p.Write(out, action))
 		ch := backend.${request}(cred, req)
 		for result := range ch {
 			if err == nil {
-				err = p.Write(conn, p.NSLCD_RESULT_BEGIN)
+				err = p.Write(out, p.NSLCD_RESULT_BEGIN)
 			}
 			if err == nil {
-				err = p.Write(conn, result)
+				err = p.Write(out, result)
 			}
 		}
 		maybePanic(err)
-		maybePanic(p.Write(conn, p.NSLCD_RESULT_END))
+		maybePanic(p.Write(out, p.NSLCD_RESULT_END))
 		return nil
 EOT
 done < "$requests"