2 files changed, 65 insertions, 36 deletions
diff --git a/cn-fileserver.service b/cn-fileserver.service
index 007585a..9ebbd7b 100644
--- a/cn-fileserver.service
+++ b/cn-fileserver.service
@@ -4,7 +4,7 @@ After=network.target
 
 [Service]
 Type=simple
-ExecStart=/usr/local/bin/cn-fileserver 0.0.0.0 /var/contnet/cnroot
+ExecStart=/usr/local/bin/cn-fileserver -listen 0.0.0.0 -dir /var/contnet/cnroot
 PrivateTmp=yes
 ProtectSystem=strict
 ProtectHome=yes
diff --git a/fileserver.go b/fileserver.go
index 5145701..e1431c4 100644
--- a/fileserver.go
+++ b/fileserver.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"flag"
 	"fmt"
 	"io"
 	"mime"
@@ -18,6 +19,12 @@ import (
 	"contnet.org/lib/cnp-go"
 )
 
+var (
+	flagSelect = flag.String("select", "byte,info,cnm", "comma-separated list of enabled selectors")
+	flagListen = flag.String("listen", "localhost:25454", "host and/or port to listen on for CNP connections")
+	flagDir    = flag.String("dir", ".", "directory to serve files from")
+)
+
 func init() {
 	mime.AddExtensionType(".cnm", "text/cnm")
 }
@@ -184,8 +191,7 @@ func getFile(dir, pth string, strict bool) (*os.File, os.FileInfo) {
 func byteSelect(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileInfo) {
 	size := stat.Size()
 	sel, selq := r.Select()
-	w.Response().SetSelect(sel, selq)
-	rej := cnp.ErrorRejected{Reason: "invalid byte selector"}
+	rej := cnp.ErrorInvalid{Reason: "invalid byte selector"}
 
 	ss := strings.Split(selq, "-")
 	if len(ss) != 2 {
@@ -209,6 +215,10 @@ func byteSelect(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileIn
 		if err != nil {
 			panic(rej)
 		}
+		if ub < ua {
+			panic(rej)
+		}
+		ub++
 		fr = io.LimitReader(fr, int64(ub-ua))
 		if int64(ub) < size {
 			size = int64(ub)
@@ -218,6 +228,7 @@ func byteSelect(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileIn
 	if size < 0 {
 		size = 0
 	}
+	w.Response().SetSelect(sel, strconv.FormatInt(int64(ua), 10)+"-"+strconv.FormatInt(int64(ua)+size-1, 10))
 	w.Response().SetLength(size)
 
 	if _, err := io.Copy(w, fr); err != nil {
@@ -231,10 +242,12 @@ func cnmSelect(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileInf
 		panic(err)
 	}
 
-	_, selq := r.Select()
+	sel, selq := r.Select()
+	w.Response().SetSelect(sel, selq)
+
 	sdoc, err := doc.Select(selq)
 	if err != nil {
-		panic(cnp.ErrorRejected{Reason: err.Error()})
+		panic(cnp.ErrorInvalid{Reason: err.Error()})
 	}
 	if sdoc == nil {
 		return
@@ -245,29 +258,45 @@ func cnmSelect(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileInf
 	}
 }
 
-func serveFile(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileInfo, typ string) {
-	switch sel, _ := r.Select(); sel {
-	case "":
-		w.Response().SetLength(stat.Size())
-		if _, err := io.Copy(w, f); err != nil {
-			panic(err)
-		}
+func infoSelect(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileInfo) {
+	sel, selq := r.Select()
+	if selq != "" {
+		panic(cnp.ErrorInvalid{Reason: "invalid info selector"})
+	}
+	resp := w.Response()
+	resp.SetLength(stat.Size())
+	var buf bytes.Buffer
+	resp.Write(&buf)
+	resp.Header = cnp.NewHeader(cnp.IntentOK, nil)
+	resp.SetLength(int64(buf.Len()))
+	resp.SetSelect(sel, selq)
+	if _, err := io.Copy(w, &buf); err != nil {
+		panic(err)
+	}
+}
+
+func noSelect(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileInfo) {
+	w.Response().SetLength(stat.Size())
+	if _, err := io.Copy(w, f); err != nil {
+		panic(err)
+	}
+}
+
+func serveFile(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileInfo, typ string, enabled map[string]bool) {
+	sel, _ := r.Select()
+
+	if !enabled[sel] {
+		noSelect(w, r, f, stat)
 		return
+	}
 
+	switch sel {
 	case "byte":
 		byteSelect(w, r, f, stat)
 		return
 
 	case "info":
-		resp := w.Response()
-		resp.SetLength(stat.Size())
-		var buf bytes.Buffer
-		resp.Write(&buf)
-		resp.Header = cnp.NewHeader(cnp.IntentOK, nil)
-		resp.SetLength(int64(buf.Len()))
-		if _, err := io.Copy(w, &buf); err != nil {
-			panic(err)
-		}
+		infoSelect(w, r, f, stat)
 		return
 
 	case "cnm":
@@ -280,7 +309,7 @@ func serveFile(w cnp.ResponseWriter, r *cnp.Request, f *os.File, stat os.FileInf
 	panic(cnp.ErrorNotSupported{Reason: "unsupported selector"})
 }
 
-func server(addr, dir string) {
+func server(addr, dir string, sel map[string]bool) {
 	fmt.Printf("listening on %q, serving dir %q\n", addr, dir)
 	panic(cnp.ListenAndServe(addr, cnp.HandlerFunc(func(w cnp.ResponseWriter, r *cnp.Request) {
 		f, stat := getFile(dir, r.Path(), false)
@@ -307,23 +336,23 @@ func server(addr, dir string) {
 		w.Response().SetName(stat.Name())
 		w.Response().SetType(typ)
 
-		serveFile(w, r, f, stat, typ)
+		serveFile(w, r, f, stat, typ, sel)
 	})))
 }
 
 func main() {
-	if len(os.Args) > 3 {
-		prog := path.Base(os.Args[0])
-		fmt.Fprintf(os.Stderr, "%s: usage: %s [HOST[:PORT]] [DIR]", prog, prog)
-		os.Exit(2)
-	}
-	addr := "localhost"
-	dir := "."
-	if len(os.Args) >= 2 {
-		addr = os.Args[1]
-	}
-	if len(os.Args) >= 3 {
-		dir = os.Args[2]
+	flag.Parse()
+
+	sel := map[string]bool{}
+	if *flagSelect != "" {
+		for _, s := range strings.Split(*flagSelect, ",") {
+			switch s {
+			case "byte", "info", "cnm":
+				sel[s] = true
+			default:
+				panic("unknown selector: " + s)
+			}
+		}
 	}
-	server(addr, dir)
+	server(*flagListen, *flagDir, sel)
 }