From 5986015358e5b7cf523122bbdb6831dccf7ca306 Mon Sep 17 00:00:00 2001 From: Horus3 Date: Mon, 11 May 2015 23:47:47 +0200 Subject: Some fixes, shows 404 page on directory not found. --- main.go | 29 ++++++----- recorder.go | 72 ++++++++++++++++++++++++++ template.go | 166 +++++++++++++++++++++++++++++------------------------------- 3 files changed, 168 insertions(+), 99 deletions(-) create mode 100644 recorder.go diff --git a/main.go b/main.go index 9ec7157..452ed35 100644 --- a/main.go +++ b/main.go @@ -10,24 +10,20 @@ import ( "strings" ) -func accessLog(h http.Handler) http.Handler { +func accessLog(h http.Handler, quiet bool) http.Handler { t := &TemplateHandler{handler: h} - //t := &TemplateHandler{} fn := func(w http.ResponseWriter, r *http.Request) { - log.Println(r.Method, r.URL.Path, r.RemoteAddr) + if !quiet { + log.Println(r.Method, r.URL.Path, r.RemoteAddr) + } t.ServeHTTP(w, r) } return http.HandlerFunc(fn) } -/* -func TemplateHandler(h http.Handler) http.Handler { - h.handler -} -*/ - -func uploadHandler(dir string) http.Handler { +func uploadHandler(dir string, quiet bool) http.Handler { fn := func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Server", "uhttpd") file, header, err := r.FormFile("file") if err != nil { @@ -60,7 +56,9 @@ func uploadHandler(dir string) http.Handler { return } - log.Println(r.Method, r.URL.Path, header.Filename, r.RemoteAddr) + if !quiet { + log.Println(r.Method, r.URL.Path, header.Filename, r.RemoteAddr) + } w.Write([]byte("Uploaded " + header.Filename)) } return http.HandlerFunc(fn) @@ -72,6 +70,7 @@ func main() { dir_f := flag.String("dir", ".", "Directory to serve.") disallow_upl_f := flag.Bool("disallow-upload", false, "Disallow uploads to /upload.") upl_dir_f := flag.String("upload-dir", ".", "Directory to store uploaded files.") + quiet_f := flag.Bool("quiet", false, "Be quiet, suppress output.") flag.Parse() @@ -80,13 +79,15 @@ func main() { port = *port_f } - fmt.Println("Starting uhttpd serving \"" + *dir_f + "\" on " + *ip_f + ":" + port + ".") + if !*quiet_f { + fmt.Println("Starting uhttpd serving \"" + *dir_f + "\" on " + *ip_f + ":" + port + ".") + } mux := http.NewServeMux() if !*disallow_upl_f { os.MkdirAll(*upl_dir_f, 0755) - mux.Handle("/upload", uploadHandler(*upl_dir_f)) + mux.Handle("/upload", uploadHandler(*upl_dir_f, *quiet_f)) } - mux.Handle("/", accessLog(http.FileServer(http.Dir(*dir_f)))) + mux.Handle("/", accessLog(http.FileServer(http.Dir(*dir_f)), *quiet_f)) log.Fatal(http.ListenAndServe(*ip_f+":"+port, mux)) } diff --git a/recorder.go b/recorder.go new file mode 100644 index 0000000..90ef1a7 --- /dev/null +++ b/recorder.go @@ -0,0 +1,72 @@ +package main + +/* + Copy from net/http/httptest + Please read there for the license. +*/ + +import ( + "bytes" + "net/http" +) + +// ResponseRecorder is an implementation of http.ResponseWriter that +// records its mutations for later inspection in tests. +type ResponseRecorder struct { + Code int // the HTTP response code from WriteHeader + HeaderMap http.Header // the HTTP response headers + Body *bytes.Buffer // if non-nil, the bytes.Buffer to append written data to + Flushed bool + + wroteHeader bool +} + +// NewRecorder returns an initialized ResponseRecorder. +func NewRecorder() *ResponseRecorder { + return &ResponseRecorder{ + HeaderMap: make(http.Header), + Body: new(bytes.Buffer), + Code: 200, + } +} + +// DefaultRemoteAddr is the default remote address to return in RemoteAddr if +// an explicit DefaultRemoteAddr isn't set on ResponseRecorder. +const DefaultRemoteAddr = "1.2.3.4" + +// Header returns the response headers. +func (rw *ResponseRecorder) Header() http.Header { + m := rw.HeaderMap + if m == nil { + m = make(http.Header) + rw.HeaderMap = m + } + return m +} + +// Write always succeeds and writes to rw.Body, if not nil. +func (rw *ResponseRecorder) Write(buf []byte) (int, error) { + if !rw.wroteHeader { + rw.WriteHeader(200) + } + if rw.Body != nil { + rw.Body.Write(buf) + } + return len(buf), nil +} + +// WriteHeader sets rw.Code. +func (rw *ResponseRecorder) WriteHeader(code int) { + if !rw.wroteHeader { + rw.Code = code + } + rw.wroteHeader = true +} + +// Flush sets rw.Flushed to true. +func (rw *ResponseRecorder) Flush() { + if !rw.wroteHeader { + rw.WriteHeader(200) + } + rw.Flushed = true +} diff --git a/template.go b/template.go index c90b13c..017690d 100644 --- a/template.go +++ b/template.go @@ -4,7 +4,6 @@ import ( "html/template" "log" "net/http" - "net/http/httptest" "strings" ) @@ -14,44 +13,69 @@ type TemplateHandler struct { func (t *TemplateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Server", "uhttpd") + // We use our custom handler to prettify directories. Files are served via default. - if strings.HasSuffix(r.URL.Path, "/") { - r.Header.Del("If-Modified-Since") - rec := httptest.NewRecorder() + if !strings.HasSuffix(r.URL.Path, "/") { + t.handler.ServeHTTP(w, r) + return + } - defer rec.Body.Reset() + rec := NewRecorder() - // passing the recorder instead of the real ResponseWriter - t.handler.ServeHTTP(rec, r) + defer rec.Body.Reset() - // we copy the original headers first - for k, v := range rec.Header() { - w.Header()[k] = v - } + // passing the recorder instead of the real ResponseWriter + t.handler.ServeHTTP(rec, r) - // BUG: Sometimes we get a 0 byte response - if rec.Body.Len() == 0 { - log.Println("Error: Empty body. Requested:", r.URL.Path) - } + // let the standard lib handle all the caching + if rec.Code > 300 && rec.Code < 400 { + log.Println("Code: ", rec.Code) + t.handler.ServeHTTP(w, r) + return + } - // we serve a file instead of a html page - if !strings.Contains(w.Header().Get("Content-Type"), "text/html") { - w.Write(rec.Body.Bytes()) + // we copy the original headers first + for k, v := range rec.Header() { + w.Header()[k] = v + } + + // not found + if rec.Code == 404 { + w.Header().Set("Content-Type", "text/html") + w.WriteHeader(404) + tmpl := template.New("404") + tmpl, err := tmpl.Parse(get404()) + if err != nil { + log.Println(err.Error()) + w.WriteHeader(500) + w.Write([]byte(err.Error())) return } - data := rec.Body.String() - - // we serve the directoy page - if strings.HasPrefix(data, "
") {
-			execTemplate(w, r, data)
-		} else {
-			w.Write(rec.Body.Bytes())
+		err = tmpl.Execute(w, struct{ URL string }{URL: r.URL.Path})
+		if err != nil {
+			log.Println(err.Error())
+			w.WriteHeader(500)
+			w.Write([]byte(err.Error()))
+			return
 		}
+		return
+	}
 
+	// we serve a file instead of a html page
+	if !strings.Contains(w.Header().Get("Content-Type"), "text/html") {
+		w.Write(rec.Body.Bytes())
 		return
+	}
+
+	data := rec.Body.String()
+
+	// we serve the directoy page
+	if strings.HasPrefix(data, "
") {
+		execTemplate(w, r, data)
 	} else {
-		t.handler.ServeHTTP(w, r)
+		w.Write(rec.Body.Bytes())
 	}
 }
 
@@ -59,10 +83,9 @@ func execTemplate(w http.ResponseWriter, r *http.Request, data string) {
 	data = strings.Replace(data, "
", "", 1)
 	data = strings.Replace(data, "
", "", 1)
 	data = strings.Replace(data, "", "
", -1)
-	data = ".. (up a dir)

" + data
 
 	tmpl := template.New("page")
-	tmpl, err := tmpl.Parse(renderTemplate())
+	tmpl, err := tmpl.Parse(getTemplate())
 	if err != nil {
 		log.Println(err.Error())
 		w.WriteHeader(500)
@@ -84,73 +107,25 @@ func NewTemplateHandler(handler http.Handler, allowedHost string) *TemplateHandl
 	return &TemplateHandler{handler: handler}
 }
 
-func renderTemplate() string {
-	html := `
+func getTemplate() string {
+	return `
 
 
 	
 	{{.URL}}
 	
-	
+	
 
 
 	

 		

 			

-				

-				
{{.URL}}

+				


+				{{.URL}}

 				

+					{{if ne .URL "/"}}
+					.. (up a dir)


+					{{end}}
 					{{.Data}}
 				

 			

@@ -171,6 +146,27 @@ div.icon div.chimney {
 
 
 `
+}
 
-	return html
+func get404() string {
+	return `
+
+
+	
+	Not Found - {{.URL}}
+	
+	
+
+
+

+	

+		
Not Found - {{.URL}}

+		

+		Please check for typos in your url. 

+		

+		Go back to root

+	

+

+
+`
 }
-- 
cgit v1.2.3