diff --git a/README.md b/README.md index 6205542..d5a1b4d 100644 --- a/README.md +++ b/README.md @@ -75,23 +75,29 @@ The configuration is pretty simple: ```ini [ application ] -listen = :8438 -application-url = http://host:port/path -clamd-url = http://host:port/ -log-file = /var/log/clammit.log -debug = true -test-pages = true +listen = :8438 +application-url = http://host:port/path +clamd-url = http://host:port/ +virus-status-code = 418 +virus-response-body = "{ \"error\": \"File Contains Virus\"}" +virus-response-content-type = "application/json" +log-file = /var/log/clammit.log +debug = true +test-pages = true ``` -Setting | Description -:---------------| :----------------------------------------------------------------------------- -listen | The listen address (see below) -clamd-url | The URL of the clamd server -application-url | (Optional) Forward all requests to this application -log-file | (Optional) The clammit log file, if ommitted will log to stdout -test-pages | (Optional) If true, clammit will also offer up a page to perform test uploads -debug | (Optional) If true, more things will be logged -debug-clam | (Optional) If true, the response from ClamAV will be logged +Setting | Description +:-----------------| :----------------------------------------------------------------------------- +listen | The listen address (see below) +clamd-url | The URL of the clamd server +application-url | (Optional) Forward all requests to this application +virus-status-code | (Optional) The HTTP status code to return when a virus is found +virus-response-body | (Optional) The HTTP body to return when a virus is found +virus-response-content-type | (Optional) The Content-Type header to return when a virus is found +log-file | (Optional) The clammit log file, if ommitted will log to stdout +test-pages | (Optional) If true, clammit will also offer up a page to perform test uploads +debug | (Optional) If true, more things will be logged +debug-clam | (Optional) If true, the response from ClamAV will be logged The listen address can be a TCP port or Unix socket, e.g.: diff --git a/src/clammit/main.go b/src/clammit/main.go index 3619b08..1de5fc1 100644 --- a/src/clammit/main.go +++ b/src/clammit/main.go @@ -67,6 +67,9 @@ type ApplicationConfig struct { ClamdURL string `gcfg:"clamd-url"` // The HTTP status code to return when a virus is found VirusStatusCode int `gcfg:"virus-status-code"` + // The HTTP status code to return when a virus is found + VirusResponseBody string `gcfg:"virus-response-body"` + VirusResponseContentType string `gcfg:"virus-response-content-type"` // If the body content-length exceeds this value, it will be written to // disk. Below it, we'll hold the whole body in memory to improve speed. ContentMemoryThreshold int64 `gcfg:"content-memory-threshold"` @@ -84,16 +87,18 @@ type ApplicationConfig struct { // Default configuration // var DefaultApplicationConfig = ApplicationConfig{ - Listen: ":8438", - SocketPerms: "0777", - ApplicationURL: "", - ClamdURL: "", - VirusStatusCode: 418, - ContentMemoryThreshold: 1024 * 1024, - Logfile: "", - TestPages: true, - Debug: false, - NumThreads: runtime.NumCPU(), + Listen: ":8438", + SocketPerms: "0777", + ApplicationURL: "", + ClamdURL: "", + VirusStatusCode: 418, + ContentMemoryThreshold: 1024 * 1024, + VirusResponseBody: "", + VirusResponseContentType: "text/plain; charset=utf-8", + Logfile: "", + TestPages: true, + Debug: false, + NumThreads: runtime.NumCPU(), } // @@ -175,8 +180,10 @@ func main() { ctx.Scanner.SetAddress(ctx.Config.App.ClamdURL) ctx.ScanInterceptor = &ScanInterceptor{ - VirusStatusCode: ctx.Config.App.VirusStatusCode, - Scanner: ctx.Scanner, + VirusStatusCode: ctx.Config.App.VirusStatusCode, + VirusResponseBody: ctx.Config.App.VirusResponseBody, + VirusResponseContentType: ctx.Config.App.VirusResponseContentType, + Scanner: ctx.Scanner, } /* diff --git a/src/clammit/scan_interceptor.go b/src/clammit/scan_interceptor.go index 8e83dae..9304468 100644 --- a/src/clammit/scan_interceptor.go +++ b/src/clammit/scan_interceptor.go @@ -19,8 +19,10 @@ import ( // The implementation of the Scan interceptor // type ScanInterceptor struct { - VirusStatusCode int - Scanner scanner.Scanner + VirusStatusCode int + VirusResponseBody string + VirusResponseContentType string + Scanner scanner.Scanner } /* @@ -113,8 +115,14 @@ func (c *ScanInterceptor) respondOnVirus(w http.ResponseWriter, filename string, http.Error(w, "Internal Server Error", 500) return true } else if hasVirus { + w.Header().Set("Content-Type", c.VirusResponseContentType) w.WriteHeader(c.VirusStatusCode) - w.Write([]byte(fmt.Sprintf("File %s has a virus!", filename))) + + if len(c.VirusResponseBody) > 0 { + w.Write([]byte(c.VirusResponseBody)) + } else { + w.Write([]byte(fmt.Sprintf("File %s has a virus!", filename))) + } return true } return false diff --git a/src/clammit/scan_interceptor_test.go b/src/clammit/scan_interceptor_test.go index d36314e..f13b6ed 100644 --- a/src/clammit/scan_interceptor_test.go +++ b/src/clammit/scan_interceptor_test.go @@ -13,6 +13,9 @@ import ( ) const virusCode = 418 +const badRequestCode = 400 +const customResponseBody = "{ \"error\": \"File Contains Virus\"}" +const customResponseContentType = "application/json" var mockVirusFound = false @@ -29,7 +32,17 @@ var scanInterceptor = ScanInterceptor{ Scanner: new(MockScanner), } +var scanInterceptorWithCustomResponse = ScanInterceptor{ + VirusStatusCode: badRequestCode, + VirusResponseBody: customResponseBody, + VirusResponseContentType: customResponseContentType, + Scanner: new(MockScanner), +} + var handler = http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { scanInterceptor.Handle(w, req, req.Body) }) +var handlerWithResponseBody = http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + scanInterceptorWithCustomResponse.Handle(w, req, req.Body) +}) func TestNonMultipartRequest_VirusFound_Without_ContentDisposition(t *testing.T) { setup() @@ -49,6 +62,29 @@ func TestNonMultipartRequest_VirusFound_Without_ContentDisposition(t *testing.T) } } +func TestNonMultipartRequest_VirusFound_With_CustomResponseFields(t *testing.T) { + setup() + mockVirusFound = true + req := newHTTPRequest("POST", "application/octet-stream", bytes.NewReader([]byte(``))) + rr := httptest.NewRecorder() + handlerWithResponseBody.ServeHTTP(rr, req) + + if status := rr.Code; status != badRequestCode { + t.Errorf("handler returned wrong status code: got %v want %v", + status, badRequestCode) + } + + if body := rr.Body.String(); body != customResponseBody { + t.Errorf("handler returned unexpected body: got %v want %v", + body, customResponseBody) + } + + if ctype := rr.Header().Get("Content-Type"); ctype != customResponseContentType { + t.Errorf("handler returned unexpected body: got %v want %v", + ctype, customResponseContentType) + } +} + func TestNonMultipartRequest_VirusFound_With_ContentDisposition(t *testing.T) { setup() mockVirusFound = true