From e0cd149b5e802db804484c7cd16da5eb5efe7753 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Thu, 4 Feb 2016 17:16:43 +0200
Subject: [PATCH] Enabled virtual hosting support in example fileserver

---
 examples/fileserver/README.md     |  1 +
 examples/fileserver/fileserver.go |  4 ++++
 fs.go                             | 32 +++++++++++++++++++++++++
 fs_test.go                        | 39 +++++++++++++++++++++++++++++++
 4 files changed, 76 insertions(+)

diff --git a/examples/fileserver/README.md b/examples/fileserver/README.md
index 5f28a45..cfc8e86 100644
--- a/examples/fileserver/README.md
+++ b/examples/fileserver/README.md
@@ -5,6 +5,7 @@
 * Supports byte range responses.
 * Generates directory index pages.
 * Supports TLS (aka SSL or HTTPS).
+* Supports virtual hosts.
 
 # How to build
 
diff --git a/examples/fileserver/fileserver.go b/examples/fileserver/fileserver.go
index 5c0fd1c..9db61f3 100644
--- a/examples/fileserver/fileserver.go
+++ b/examples/fileserver/fileserver.go
@@ -17,6 +17,7 @@ var (
 	dir                = flag.String("dir", "/usr/share/nginx/html", "Directory to serve static files from")
 	generateIndexPages = flag.Bool("generateIndexPages", true, "Whether to generate directory index pages")
 	keyFile            = flag.String("keyFile", "./ssl-cert-snakeoil.key", "Path to TLS key file")
+	vhost              = flag.Bool("vhost", false, "Enables virtual hosting by prepending the requested path with the requested hostname")
 )
 
 func main() {
@@ -29,6 +30,9 @@ func main() {
 		Compress:           *compress,
 		AcceptByteRange:    *byteRange,
 	}
+	if *vhost {
+		fs.PathRewrite = fasthttp.NewVHostPathRewriter(0)
+	}
 	h := fs.NewRequestHandler()
 
 	// Start HTTP server.
diff --git a/fs.go b/fs.go
index 0519926..765f30b 100644
--- a/fs.go
+++ b/fs.go
@@ -83,6 +83,38 @@ var (
 // The returned path may refer to ctx members. For example, ctx.Path().
 type PathRewriteFunc func(ctx *RequestCtx) []byte
 
+// NewVHostPathRewriter returns path rewriter, which strips slashesCount
+// leading slashes from the path and prepends the path with request's host,
+// thus simplifying virtual hosting for static files.
+//
+// Examples:
+//
+//   * host=foobar.com, slashesCount=0, original path="/foo/bar".
+//     Resulting path: "/foobar.com/foo/bar"
+//
+//   * host=img.aaa.com, slashesCount=1, original path="/images/123/456.jpg"
+//     Resulting path: "/img.aaa.com/123/456.jpg"
+//
+func NewVHostPathRewriter(slashesCount int) PathRewriteFunc {
+	return func(ctx *RequestCtx) []byte {
+		path := stripLeadingSlashes(ctx.Path(), slashesCount)
+		host := ctx.Host()
+		if n := bytes.IndexByte(host, '/'); n >= 0 {
+			host = nil
+		}
+		if len(host) == 0 {
+			host = strInvalidHost
+		}
+		newPath := make([]byte, len(path)+len(host)+1)
+		newPath[0] = '/'
+		copy(newPath[1:], host)
+		copy(newPath[1+len(host):], path)
+		return newPath
+	}
+}
+
+var strInvalidHost = []byte("invalid-host")
+
 // NewPathSlashesStripper returns path rewriter, which strips slashesCount
 // leading slashes from the path.
 //
diff --git a/fs_test.go b/fs_test.go
index d138c60..3aa856b 100644
--- a/fs_test.go
+++ b/fs_test.go
@@ -12,6 +12,45 @@ import (
 	"time"
 )
 
+func TestNewVHostPathRewriter(t *testing.T) {
+	var ctx RequestCtx
+	var req Request
+	req.Header.SetHost("foobar.com")
+	req.SetRequestURI("/foo/bar/baz")
+	ctx.Init(&req, nil, nil)
+
+	f := NewVHostPathRewriter(0)
+	path := f(&ctx)
+	expectedPath := "/foobar.com/foo/bar/baz"
+	if string(path) != expectedPath {
+		t.Fatalf("unexpected path %q. Expecting %q", path, expectedPath)
+	}
+
+	ctx.Request.Reset()
+	ctx.Request.SetRequestURI("https://aaa.bbb.cc/one/two/three/four?asdf=dsf")
+	f = NewVHostPathRewriter(2)
+	path = f(&ctx)
+	expectedPath = "/aaa.bbb.cc/three/four"
+	if string(path) != expectedPath {
+		t.Fatalf("unexpected path %q. Expecting %q", path, expectedPath)
+	}
+}
+
+func TestNewVHostPathRewriterMaliciousHost(t *testing.T) {
+	var ctx RequestCtx
+	var req Request
+	req.Header.SetHost("/../../../etc/passwd")
+	req.SetRequestURI("/foo/bar/baz")
+	ctx.Init(&req, nil, nil)
+
+	f := NewVHostPathRewriter(0)
+	path := f(&ctx)
+	expectedPath := "/invalid-host/foo/bar/baz"
+	if string(path) != expectedPath {
+		t.Fatalf("unexpected path %q. Expecting %q", path, expectedPath)
+	}
+}
+
 func TestServeFileHead(t *testing.T) {
 	var ctx RequestCtx
 	var req Request