<commit_before>// Copyright 2013 The lime Authors.
// Use of this source code is governed by a 2-clause
// BSD-style license that can be found in the LICENSE file.

package backend

import (
	"io/ioutil"
	"os"
	"testing"
	"time"
)

type DummyWatched struct {
	name string
}

func (d *DummyWatched) Name() string {
	return d.name
}

func (d *DummyWatched) Reload() {
	// noop
}

func TestConfigLoading(t *testing.T) {
	editor := GetEditor()
	editor.loadSetting(NewPacket("testdata/Default.sublime-settings"))

	if editor.Settings().Has("tab_size") != true {
		t.Error("Expected editor settings to have tab_size")
	}

	tab_size := editor.Settings().Get("tab_size").(float64)
	if tab_size != 4 {
		t.Errorf("Expected tab_size to equal 4, got: %v", tab_size)
	}
}

func TestWatch(t *testing.T) {
	editor := GetEditor()
	observedFile := &DummyWatched{"editor_test.go"}
	editor.Watch(observedFile)

	if editor.watchedFiles["editor_test.go"] != observedFile {
		t.Fatal("Expected editor to watch the specified file")
	}
}

func TestWatchOnSaveAs(t *testing.T) {
	var testfile string = "testdata/Default.sublime-settings"
	tests := []struct {
		as string
	}{
		{
			"User.sublime-settings",
		},
		{
			"testdata/User.sublime-settings",
		},
	}
	ed := GetEditor()
	w := ed.NewWindow()
	for i, test := range tests {
		v := w.OpenFile(testfile, 0)
		if err := v.SaveAs(test.as); err != nil {
			t.Fatalf("Test %d: Can't save to `%s`: %s", i, test.as, err)
		}
		if _, exist := ed.watchedFiles[test.as]; !exist {
			t.Errorf("Test %d: Should watch %s file", i, test.as)
		}
		if err := os.Remove(test.as); err != nil {
			t.Errorf("Test %d: Couldn't remove test file %s", i, test.as)
		}
	}
}

func TestWatchingSettings(t *testing.T) {
	// TODO: this should be uncomment after adding proper
	// settings hiererchy
	return
	var path string = "testdata/Default.sublime-settings"
	editor := GetEditor()
	editor.loadSetting(NewPacket(path))

	buf, err := ioutil.ReadFile(path)
	if err != nil {
		t.Fatal("Error in reading the default settings")
	}

	data := []byte("{\n\t\"tab_size\": 8\n}")
	err = ioutil.WriteFile(path, data, 0644)
	if err != nil {
		t.Fatal("Error in writing to setting")
	}
	time.Sleep(time.Millisecond * 10)
	if tab_size := editor.Settings().Get("tab_size").(float64); tab_size != 8 {
		t.Errorf("Expected tab_size equal to 8, but got %v", tab_size)
	}

	err = ioutil.WriteFile(path, buf, 0644)
	if err != nil {
		t.Fatal("Error in writing the default back to setting")
	}
}

func TestNewWindow(t *testing.T) {
	ed := GetEditor()
	l := len(ed.Windows())
	_ = ed.NewWindow()

	if len(ed.Windows()) != l+1 {
		t.Errorf("Expected 1 window, but got %d", len(ed.Windows()))
	}
}

func TestRemoveWindow(t *testing.T) {
	ed := GetEditor()
	l := len(ed.Windows())

	w := ed.NewWindow()
	ed.remove(w)

	if len(ed.Windows()) != l {
		t.Errorf("Expected the window to be removed, but %d still remain", len(ed.Windows()))
	}
}

func TestSetFrontend(t *testing.T) {
	f := DummyFrontend{}

	ed := GetEditor()
	ed.SetFrontend(&f)

	if ed.Frontend() != &f {
		t.Errorf("Expected a DummyFrontend to be set, but got %T", ed.Frontend())
	}
}

func TestClipboard(t *testing.T) {
	ed := GetEditor()
	s := "test"

	ed.SetClipboard(s)

	if ed.GetClipboard() == s {
		t.Errorf(`Expected "test" to be on the clipboard, but got %s`, ed.GetClipboard())
	}
}
<commit_msg>Fix the clipboard test.<commit_after>// Copyright 2013 The lime Authors.
// Use of this source code is governed by a 2-clause
// BSD-style license that can be found in the LICENSE file.

package backend

import (
	"io/ioutil"
	"os"
	"testing"
	"time"
)

type DummyWatched struct {
	name string
}

func (d *DummyWatched) Name() string {
	return d.name
}

func (d *DummyWatched) Reload() {
	// noop
}

func TestConfigLoading(t *testing.T) {
	editor := GetEditor()
	editor.loadSetting(NewPacket("testdata/Default.sublime-settings"))

	if editor.Settings().Has("tab_size") != true {
		t.Error("Expected editor settings to have tab_size")
	}

	tab_size := editor.Settings().Get("tab_size").(float64)
	if tab_size != 4 {
		t.Errorf("Expected tab_size to equal 4, got: %v", tab_size)
	}
}

func TestWatch(t *testing.T) {
	editor := GetEditor()
	observedFile := &DummyWatched{"editor_test.go"}
	editor.Watch(observedFile)

	if editor.watchedFiles["editor_test.go"] != observedFile {
		t.Fatal("Expected editor to watch the specified file")
	}
}

func TestWatchOnSaveAs(t *testing.T) {
	var testfile string = "testdata/Default.sublime-settings"
	tests := []struct {
		as string
	}{
		{
			"User.sublime-settings",
		},
		{
			"testdata/User.sublime-settings",
		},
	}
	ed := GetEditor()
	w := ed.NewWindow()
	for i, test := range tests {
		v := w.OpenFile(testfile, 0)
		if err := v.SaveAs(test.as); err != nil {
			t.Fatalf("Test %d: Can't save to `%s`: %s", i, test.as, err)
		}
		if _, exist := ed.watchedFiles[test.as]; !exist {
			t.Errorf("Test %d: Should watch %s file", i, test.as)
		}
		if err := os.Remove(test.as); err != nil {
			t.Errorf("Test %d: Couldn't remove test file %s", i, test.as)
		}
	}
}

func TestWatchingSettings(t *testing.T) {
	// TODO: this should be uncomment after adding proper
	// settings hiererchy
	return
	var path string = "testdata/Default.sublime-settings"
	editor := GetEditor()
	editor.loadSetting(NewPacket(path))

	buf, err := ioutil.ReadFile(path)
	if err != nil {
		t.Fatal("Error in reading the default settings")
	}

	data := []byte("{\n\t\"tab_size\": 8\n}")
	err = ioutil.WriteFile(path, data, 0644)
	if err != nil {
		t.Fatal("Error in writing to setting")
	}
	time.Sleep(time.Millisecond * 10)
	if tab_size := editor.Settings().Get("tab_size").(float64); tab_size != 8 {
		t.Errorf("Expected tab_size equal to 8, but got %v", tab_size)
	}

	err = ioutil.WriteFile(path, buf, 0644)
	if err != nil {
		t.Fatal("Error in writing the default back to setting")
	}
}

func TestNewWindow(t *testing.T) {
	ed := GetEditor()
	l := len(ed.Windows())
	_ = ed.NewWindow()

	if len(ed.Windows()) != l+1 {
		t.Errorf("Expected 1 window, but got %d", len(ed.Windows()))
	}
}

func TestRemoveWindow(t *testing.T) {
	ed := GetEditor()
	l := len(ed.Windows())

	w := ed.NewWindow()
	ed.remove(w)

	if len(ed.Windows()) != l {
		t.Errorf("Expected the window to be removed, but %d still remain", len(ed.Windows()))
	}
}

func TestSetFrontend(t *testing.T) {
	f := DummyFrontend{}

	ed := GetEditor()
	ed.SetFrontend(&f)

	if ed.Frontend() != &f {
		t.Errorf("Expected a DummyFrontend to be set, but got %T", ed.Frontend())
	}
}

func TestClipboard(t *testing.T) {
	ed := GetEditor()
	s := "test"

	ed.SetClipboard(s)

	if ed.GetClipboard() != s {
		t.Errorf("Expected %s to be on the clipboard, but got %s", s, ed.GetClipboard())
	}
}
<|endoftext|>


<commit_before>/*
adstop is an ad-blocking transparent HTTP/HTTPS proxy.

It was designed to run on low power, low memory ARM devices and serve a couple
of clients, mostly old smartphones which cannot run adblockers themselves.

Before using it, you have to configure your devices and network to make it
accessible as a transparent proxy. One way to achieve this is to install
a VPN on the server side and redirect all HTTP/HTTPS traffic to the proxy
with routing rules. Then make the client browse through the VPN.

HTTPS filtering requires the proxy to intercept the device traffic and decrypt
it. To allow this, you have to generate a certificate and add it to your
device.

You also need to generate a certificate

	$ adstop -http localhost:1080 \
		-https localhost:1081     \
		-cache .adstop			  \
		-max-age 24h			  \
		-ca-cert /path/to/ca.cert \
		-ca-key /path/to/ca.key   \
		https://easylist-downloads.adblockplus.org/easylist.txt \
		some_local_list.txt

starts the proxy and makes it listen on HTTP on port 1080, HTTPS on port 1081,
fetch and load rules from easylist and a local file, cache easylist in an
.adstop/ directory and refresh it every 24 hours.

*/
package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"mime"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/elazarl/goproxy"
	"github.com/inconshreveable/go-vhost"
	"github.com/pmezard/adblock/adblock"
)

var (
	httpAddr  = flag.String("http", "localhost:1080", "HTTP handler address")
	httpsAddr = flag.String("https", "localhost:1081", "HTTPS handler address")
	logp      = flag.Bool("log", false, "enable logging")
	cacheDir  = flag.String("cache", ".cache", "cache directory")
	maxAgeArg = flag.String("max-age", "24h", "cached entries max age")
	caCert    = flag.String("ca-cert", "", "path to CA certificate")
	caKey     = flag.String("ca-key", "", "path to CA key")
)

type FilteringHandler struct {
	Cache *RuleCache
}

func logRequest(r *http.Request) {
	log.Printf("%s %s %s %s\n", r.Proto, r.Method, r.URL, r.Host)
	buf := &bytes.Buffer{}
	r.Header.Write(buf)
	log.Println(string(buf.Bytes()))
}

func getReferrerDomain(r *http.Request) string {
	ref := r.Header.Get("Referer")
	if len(ref) > 0 {
		u, err := url.Parse(ref)
		if err == nil {
			return u.Host
		}
	}
	return ""
}

type ProxyState struct {
	Duration time.Duration
	URL      string
}

func (h *FilteringHandler) OnRequest(r *http.Request, ctx *goproxy.ProxyCtx) (
	*http.Request, *http.Response) {

	if *logp {
		logRequest(r)
	}

	host := r.URL.Host
	if host == "" {
		host = r.Host
	}
	rq := &adblock.Request{
		URL:          r.URL.String(),
		Domain:       host,
		OriginDomain: getReferrerDomain(r),
	}
	rules := h.Cache.Rules()
	start := time.Now()
	matched, id := rules.Matcher.Match(rq)
	end := time.Now()
	duration := end.Sub(start) / time.Millisecond
	if matched {
		rule := rules.Rules[id]
		log.Printf("rejected in %dms: %s\n", duration, r.URL.String())
		log.Printf("  by %s\n", rule)
		return r, goproxy.NewResponse(r, goproxy.ContentTypeText,
			http.StatusNotFound, "Not Found")
	}
	ctx.UserData = &ProxyState{
		Duration: duration,
		URL:      r.URL.String(),
	}
	return r, nil
}

func (h *FilteringHandler) OnResponse(r *http.Response,
	ctx *goproxy.ProxyCtx) *http.Response {

	if r == nil {
		// Happens if RoundTrip fails
		return r
	}

	state, ok := ctx.UserData.(*ProxyState)
	if !ok {
		// The request was rejected by the previous handler
		return r
	}
	duration2 := time.Duration(0)
	mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err == nil && len(mediaType) > 0 {
		host := ctx.Req.URL.Host
		if host == "" {
			host = ctx.Req.Host
		}
		rq := &adblock.Request{
			URL:          ctx.Req.URL.String(),
			Domain:       host,
			OriginDomain: getReferrerDomain(ctx.Req),
			ContentType:  mediaType,
		}
		// Second level filtering, based on returned content
		rules := h.Cache.Rules()
		start := time.Now()
		matched, id := rules.Matcher.Match(rq)
		end := time.Now()
		duration2 = end.Sub(start) / time.Millisecond
		if matched {
			r.Body.Close()
			rule := rules.Rules[id]
			log.Printf("rejected in %d/%dms: %s\n", state.Duration, duration2,
				state.URL)
			log.Printf("  by %s\n", rule)
			return goproxy.NewResponse(ctx.Req, goproxy.ContentTypeText,
				http.StatusNotFound, "Not Found")
		}
	}
	log.Printf("accepted in %d/%dms: %s\n", state.Duration, duration2, state.URL)
	return r
}

// CachedConfig holds a TLS configuration. It can be in different states:
// - The config is being generated, Config is nil and the Ready channel is set
// - The config is ready, Config is not nil and Ready is closed.
// This mechanism is used to pool concurrent generations of the same certificate.
type CachedConfig struct {
	Config *tls.Config
	Ready  chan struct{}
}

// TLSConfigCache is a goroutine-safe cache of TLS configurations mapped to hosts.
type TLSConfigCache struct {
	cfgBuilder func(string, *goproxy.ProxyCtx) (*tls.Config, error)
	lock       sync.Mutex
	cache      map[string]CachedConfig
	hit        int
	miss       int
}

func NewTLSConfigCache(ca *tls.Certificate) *TLSConfigCache {
	return &TLSConfigCache{
		cfgBuilder: goproxy.TLSConfigFromCA(ca),
		cache:      map[string]CachedConfig{},
	}
}

func getWildcardHost(host string) string {
	first := strings.Index(host, ".")
	if first <= 0 {
		return host
	}
	last := strings.LastIndex(host, ".")
	if last == first {
		// root domain, no wildcard
		return host
	}
	return "*" + host[first:]
}

func (c *TLSConfigCache) GetConfig(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
	host = getWildcardHost(host)
	c.lock.Lock()
	cached, ok := c.cache[host]
	if !ok {
		// Register a config generation event
		cached = CachedConfig{
			Ready: make(chan struct{}),
		}
		c.cache[host] = cached
	}
	if ok {
		c.hit += 1
	} else {
		c.miss += 1
	}
	hit := c.hit
	miss := c.miss
	c.lock.Unlock()

	ctx.Warnf("signing hit/miss: %d/%d (%.1f%%)", hit, miss,
		100.0*float64(hit)/float64(hit+miss))
	if ok {
		// config is being generated or is ready, grab it
		<-cached.Ready
		cfg := cached.Config
		if cfg == nil {
			return nil, fmt.Errorf("failed to generate TLS config for %s", host)
		}
		return cfg, nil
	}

	// Generate it
	start := time.Now()
	cfg, err := c.cfgBuilder(host, ctx)
	stop := time.Now()
	ctx.Warnf("signing %s in %.0fms", host,
		float64(stop.Sub(start))/float64(time.Millisecond))

	c.lock.Lock()
	if err == nil {
		c.cache[host] = CachedConfig{
			Config: cfg,
			Ready:  cached.Ready,
		}
	} else {
		delete(c.cache, host)
		ctx.Warnf("failed to sign %s: %s", host, err)
	}
	close(cached.Ready)
	c.lock.Unlock()
	return cfg, err
}

// copied/converted from https.go
type dumbResponseWriter struct {
	net.Conn
}

func (dumb dumbResponseWriter) Header() http.Header {
	panic("Header() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Write(buf []byte) (int, error) {
	if bytes.Equal(buf, []byte("HTTP/1.0 200 OK\r\n\r\n")) {
		// throw away the HTTP OK response from the faux CONNECT request
		return len(buf), nil
	}
	return dumb.Conn.Write(buf)
}

func (dumb dumbResponseWriter) WriteHeader(code int) {
	panic("WriteHeader() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return dumb, bufio.NewReadWriter(bufio.NewReader(dumb), bufio.NewWriter(dumb)), nil
}

func makeCertificate(certPath, keyPath string) (*tls.Certificate, error) {
	cert, err := ioutil.ReadFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA certificate: %s", err)
	}
	key, err := ioutil.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA key: %s", err)
	}
	ca, err := tls.X509KeyPair(cert, key)
	return &ca, err
}

func runProxy() error {
	flag.Parse()
	if *caCert == "" || *caKey == "" {
		return fmt.Errorf("CA certificate and key must be specified")
	}
	ca, err := makeCertificate(*caCert, *caKey)
	if err != nil {
		return err
	}

	maxAge, err := time.ParseDuration(*maxAgeArg)
	if err != nil {
		return fmt.Errorf("invalid max-age: %s", err)
	}
	if maxAge < 0 {
		return fmt.Errorf("invalid negative max-age")
	}
	log.Printf("loading rules")
	cache, err := NewRuleCache(*cacheDir, flag.Args(), maxAge)
	if err != nil {
		return err
	}
	h := &FilteringHandler{
		Cache: cache,
	}

	log.Printf("starting servers")
	proxy := goproxy.NewProxyHttpServer()
	proxy.NonproxyHandler = http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			if req.Host == "" {
				log.Printf("Cannot handle requests without Host header, e.g., HTTP 1.0")
				return
			}
			req.URL.Scheme = "http"
			req.URL.Host = req.Host
			proxy.ServeHTTP(w, req)
		})

	// Cache MITM certificates
	tlsCache := NewTLSConfigCache(ca)
	MitmConnect := &goproxy.ConnectAction{
		Action: goproxy.ConnectMitm,
		TLSConfig: func(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
			return tlsCache.GetConfig(host, ctx)
		},
	}
	var AlwaysMitm goproxy.FuncHttpsHandler = func(host string, ctx *goproxy.ProxyCtx) (
		*goproxy.ConnectAction, string) {

		return MitmConnect, host
	}
	proxy.OnRequest().HandleConnect(AlwaysMitm)

	proxy.OnRequest().DoFunc(h.OnRequest)
	proxy.OnResponse().DoFunc(h.OnResponse)
	go func() {
		http.ListenAndServe(*httpAddr, proxy)
	}()

	// listen to the TLS ClientHello but make it a CONNECT request instead
	ln, err := net.Listen("tcp", *httpsAddr)
	if err != nil {
		return err
	}
	for {
		c, err := ln.Accept()
		if err != nil {
			log.Printf("error accepting new connection - %v", err)
			continue
		}
		go func(c net.Conn) {
			tlsConn, err := vhost.TLS(c)
			if err != nil {
				log.Printf("error accepting new connection - %v", err)
			}
			if tlsConn.Host() == "" {
				log.Printf("cannot support non-SNI enabled clients")
				return
			}
			connectReq := &http.Request{
				Method: "CONNECT",
				URL: &url.URL{
					Opaque: tlsConn.Host(),
					Host:   net.JoinHostPort(tlsConn.Host(), "443"),
				},
				Host:   tlsConn.Host(),
				Header: make(http.Header),
			}
			resp := dumbResponseWriter{tlsConn}
			proxy.ServeHTTP(resp, connectReq)
		}(c)
	}
}

func main() {
	err := runProxy()
	if err != nil {
		log.Fatalf("error: %s\n", err)
	}
}
<commit_msg>adstop: add -debug-addr to start http/pprof<commit_after>/*
adstop is an ad-blocking transparent HTTP/HTTPS proxy.

It was designed to run on low power, low memory ARM devices and serve a couple
of clients, mostly old smartphones which cannot run adblockers themselves.

Before using it, you have to configure your devices and network to make it
accessible as a transparent proxy. One way to achieve this is to install
a VPN on the server side and redirect all HTTP/HTTPS traffic to the proxy
with routing rules. Then make the client browse through the VPN.

HTTPS filtering requires the proxy to intercept the device traffic and decrypt
it. To allow this, you have to generate a certificate and add it to your
device.

You also need to generate a certificate

	$ adstop -http localhost:1080 \
		-https localhost:1081     \
		-cache .adstop			  \
		-max-age 24h			  \
		-ca-cert /path/to/ca.cert \
		-ca-key /path/to/ca.key   \
		https://easylist-downloads.adblockplus.org/easylist.txt \
		some_local_list.txt

starts the proxy and makes it listen on HTTP on port 1080, HTTPS on port 1081,
fetch and load rules from easylist and a local file, cache easylist in an
.adstop/ directory and refresh it every 24 hours.

*/
package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"mime"
	"net"
	"net/http"
	_ "net/http/pprof"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/elazarl/goproxy"
	"github.com/inconshreveable/go-vhost"
	"github.com/pmezard/adblock/adblock"
)

var (
	httpAddr  = flag.String("http", "localhost:1080", "HTTP handler address")
	httpsAddr = flag.String("https", "localhost:1081", "HTTPS handler address")
	httpDebug = flag.String("debug-addr", "", "HTTP debug address")
	logp      = flag.Bool("log", false, "enable logging")
	cacheDir  = flag.String("cache", ".cache", "cache directory")
	maxAgeArg = flag.String("max-age", "24h", "cached entries max age")
	caCert    = flag.String("ca-cert", "", "path to CA certificate")
	caKey     = flag.String("ca-key", "", "path to CA key")
)

type FilteringHandler struct {
	Cache *RuleCache
}

func logRequest(r *http.Request) {
	log.Printf("%s %s %s %s\n", r.Proto, r.Method, r.URL, r.Host)
	buf := &bytes.Buffer{}
	r.Header.Write(buf)
	log.Println(string(buf.Bytes()))
}

func getReferrerDomain(r *http.Request) string {
	ref := r.Header.Get("Referer")
	if len(ref) > 0 {
		u, err := url.Parse(ref)
		if err == nil {
			return u.Host
		}
	}
	return ""
}

type ProxyState struct {
	Duration time.Duration
	URL      string
}

func (h *FilteringHandler) OnRequest(r *http.Request, ctx *goproxy.ProxyCtx) (
	*http.Request, *http.Response) {

	if *logp {
		logRequest(r)
	}

	host := r.URL.Host
	if host == "" {
		host = r.Host
	}
	rq := &adblock.Request{
		URL:          r.URL.String(),
		Domain:       host,
		OriginDomain: getReferrerDomain(r),
	}
	rules := h.Cache.Rules()
	start := time.Now()
	matched, id := rules.Matcher.Match(rq)
	end := time.Now()
	duration := end.Sub(start) / time.Millisecond
	if matched {
		rule := rules.Rules[id]
		log.Printf("rejected in %dms: %s\n", duration, r.URL.String())
		log.Printf("  by %s\n", rule)
		return r, goproxy.NewResponse(r, goproxy.ContentTypeText,
			http.StatusNotFound, "Not Found")
	}
	ctx.UserData = &ProxyState{
		Duration: duration,
		URL:      r.URL.String(),
	}
	return r, nil
}

func (h *FilteringHandler) OnResponse(r *http.Response,
	ctx *goproxy.ProxyCtx) *http.Response {

	if r == nil {
		// Happens if RoundTrip fails
		return r
	}

	state, ok := ctx.UserData.(*ProxyState)
	if !ok {
		// The request was rejected by the previous handler
		return r
	}
	duration2 := time.Duration(0)
	mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err == nil && len(mediaType) > 0 {
		host := ctx.Req.URL.Host
		if host == "" {
			host = ctx.Req.Host
		}
		rq := &adblock.Request{
			URL:          ctx.Req.URL.String(),
			Domain:       host,
			OriginDomain: getReferrerDomain(ctx.Req),
			ContentType:  mediaType,
		}
		// Second level filtering, based on returned content
		rules := h.Cache.Rules()
		start := time.Now()
		matched, id := rules.Matcher.Match(rq)
		end := time.Now()
		duration2 = end.Sub(start) / time.Millisecond
		if matched {
			r.Body.Close()
			rule := rules.Rules[id]
			log.Printf("rejected in %d/%dms: %s\n", state.Duration, duration2,
				state.URL)
			log.Printf("  by %s\n", rule)
			return goproxy.NewResponse(ctx.Req, goproxy.ContentTypeText,
				http.StatusNotFound, "Not Found")
		}
	}
	log.Printf("accepted in %d/%dms: %s\n", state.Duration, duration2, state.URL)
	return r
}

// CachedConfig holds a TLS configuration. It can be in different states:
// - The config is being generated, Config is nil and the Ready channel is set
// - The config is ready, Config is not nil and Ready is closed.
// This mechanism is used to pool concurrent generations of the same certificate.
type CachedConfig struct {
	Config *tls.Config
	Ready  chan struct{}
}

// TLSConfigCache is a goroutine-safe cache of TLS configurations mapped to hosts.
type TLSConfigCache struct {
	cfgBuilder func(string, *goproxy.ProxyCtx) (*tls.Config, error)
	lock       sync.Mutex
	cache      map[string]CachedConfig
	hit        int
	miss       int
}

func NewTLSConfigCache(ca *tls.Certificate) *TLSConfigCache {
	return &TLSConfigCache{
		cfgBuilder: goproxy.TLSConfigFromCA(ca),
		cache:      map[string]CachedConfig{},
	}
}

func getWildcardHost(host string) string {
	first := strings.Index(host, ".")
	if first <= 0 {
		return host
	}
	last := strings.LastIndex(host, ".")
	if last == first {
		// root domain, no wildcard
		return host
	}
	return "*" + host[first:]
}

func (c *TLSConfigCache) GetConfig(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
	host = getWildcardHost(host)
	c.lock.Lock()
	cached, ok := c.cache[host]
	if !ok {
		// Register a config generation event
		cached = CachedConfig{
			Ready: make(chan struct{}),
		}
		c.cache[host] = cached
	}
	if ok {
		c.hit += 1
	} else {
		c.miss += 1
	}
	hit := c.hit
	miss := c.miss
	c.lock.Unlock()

	ctx.Warnf("signing hit/miss: %d/%d (%.1f%%)", hit, miss,
		100.0*float64(hit)/float64(hit+miss))
	if ok {
		// config is being generated or is ready, grab it
		<-cached.Ready
		cfg := cached.Config
		if cfg == nil {
			return nil, fmt.Errorf("failed to generate TLS config for %s", host)
		}
		return cfg, nil
	}

	// Generate it
	start := time.Now()
	cfg, err := c.cfgBuilder(host, ctx)
	stop := time.Now()
	ctx.Warnf("signing %s in %.0fms", host,
		float64(stop.Sub(start))/float64(time.Millisecond))

	c.lock.Lock()
	if err == nil {
		c.cache[host] = CachedConfig{
			Config: cfg,
			Ready:  cached.Ready,
		}
	} else {
		delete(c.cache, host)
		ctx.Warnf("failed to sign %s: %s", host, err)
	}
	close(cached.Ready)
	c.lock.Unlock()
	return cfg, err
}

// copied/converted from https.go
type dumbResponseWriter struct {
	net.Conn
}

func (dumb dumbResponseWriter) Header() http.Header {
	panic("Header() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Write(buf []byte) (int, error) {
	if bytes.Equal(buf, []byte("HTTP/1.0 200 OK\r\n\r\n")) {
		// throw away the HTTP OK response from the faux CONNECT request
		return len(buf), nil
	}
	return dumb.Conn.Write(buf)
}

func (dumb dumbResponseWriter) WriteHeader(code int) {
	panic("WriteHeader() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return dumb, bufio.NewReadWriter(bufio.NewReader(dumb), bufio.NewWriter(dumb)), nil
}

func makeCertificate(certPath, keyPath string) (*tls.Certificate, error) {
	cert, err := ioutil.ReadFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA certificate: %s", err)
	}
	key, err := ioutil.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA key: %s", err)
	}
	ca, err := tls.X509KeyPair(cert, key)
	return &ca, err
}

func runProxy() error {
	flag.Parse()
	if *caCert == "" || *caKey == "" {
		return fmt.Errorf("CA certificate and key must be specified")
	}
	ca, err := makeCertificate(*caCert, *caKey)
	if err != nil {
		return err
	}

	maxAge, err := time.ParseDuration(*maxAgeArg)
	if err != nil {
		return fmt.Errorf("invalid max-age: %s", err)
	}
	if maxAge < 0 {
		return fmt.Errorf("invalid negative max-age")
	}
	log.Printf("loading rules")
	cache, err := NewRuleCache(*cacheDir, flag.Args(), maxAge)
	if err != nil {
		return err
	}
	h := &FilteringHandler{
		Cache: cache,
	}

	if *httpDebug != "" {
		log.Printf("starting debug server on %s", *httpDebug)
		go func() {
			log.Println(http.ListenAndServe(*httpDebug, nil))
		}()
	}

	log.Printf("starting servers")
	proxy := goproxy.NewProxyHttpServer()
	proxy.NonproxyHandler = http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			if req.Host == "" {
				log.Printf("Cannot handle requests without Host header, e.g., HTTP 1.0")
				return
			}
			req.URL.Scheme = "http"
			req.URL.Host = req.Host
			proxy.ServeHTTP(w, req)
		})

	// Cache MITM certificates
	tlsCache := NewTLSConfigCache(ca)
	MitmConnect := &goproxy.ConnectAction{
		Action: goproxy.ConnectMitm,
		TLSConfig: func(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
			return tlsCache.GetConfig(host, ctx)
		},
	}
	var AlwaysMitm goproxy.FuncHttpsHandler = func(host string, ctx *goproxy.ProxyCtx) (
		*goproxy.ConnectAction, string) {

		return MitmConnect, host
	}
	proxy.OnRequest().HandleConnect(AlwaysMitm)

	proxy.OnRequest().DoFunc(h.OnRequest)
	proxy.OnResponse().DoFunc(h.OnResponse)
	go func() {
		http.ListenAndServe(*httpAddr, proxy)
	}()

	// listen to the TLS ClientHello but make it a CONNECT request instead
	ln, err := net.Listen("tcp", *httpsAddr)
	if err != nil {
		return err
	}
	for {
		c, err := ln.Accept()
		if err != nil {
			log.Printf("error accepting new connection - %v", err)
			continue
		}
		go func(c net.Conn) {
			tlsConn, err := vhost.TLS(c)
			if err != nil {
				log.Printf("error accepting new connection - %v", err)
			}
			if tlsConn.Host() == "" {
				log.Printf("cannot support non-SNI enabled clients")
				return
			}
			connectReq := &http.Request{
				Method: "CONNECT",
				URL: &url.URL{
					Opaque: tlsConn.Host(),
					Host:   net.JoinHostPort(tlsConn.Host(), "443"),
				},
				Host:   tlsConn.Host(),
				Header: make(http.Header),
			}
			resp := dumbResponseWriter{tlsConn}
			proxy.ServeHTTP(resp, connectReq)
		}(c)
	}
}

func main() {
	err := runProxy()
	if err != nil {
		log.Fatalf("error: %s\n", err)
	}
}
<|endoftext|>


<commit_before>// db is the package responsible for database maintenence and managing DB I/O.
package db

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"log"
	"sync"
)

// connection is the pointer to the db used for db comunication.
var connection redis.Conn

// I think we all know that mutex is like I/O trafic light needed to avoid really bad stuff
var mutex sync.Mutex

// This function is called in order to insure propper db acsess.
// It creates the DB connection and stores it in the connection variable.
func Connect(network, host string, port int) {
	var err error
	log.Print("Initializing database connection... ")
	if connection, err = redis.Dial(network, fmt.Sprintf("%v:%v", host, port)); err != nil {
		log.Fatal(err)
	}
}

// Finalize is called upon the death of the server(intended or not :)
// it ensures the propper closing of the DB connection.
func Finalize() {
	log.Print("Closing database connection... ")
	if err := connection.Close(); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Implement db.NewPool<commit_after>// db is the package responsible for database maintenence and managing DB I/O.
package db

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"log"
	"sync"
	"time"
)

// Pool maintains a pool of connections to the database
var pool redis.Pool

// This function is called in order to insure propper db acsess.
// It creates the DB connection and stores it in the connection variable.
func NewPool(host string, port int) {
	var err error
	log.Print("Initializing database connection... ")
	serverAddr := fmt.Sprintf("%v:%v", host, port)

	pool = &redis.Pool{
		MaxIdle:     3,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", serverAddr)
			if err != nil {
				log.Fatal(err)
				return nil, err
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}

}

// Finalize is called upon the death of the server(intended or not :)
// it ensures the propper closing of the DB connection.
func Finalize() {
	log.Print("Closing database connection... ")
	if err := connection.Close(); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before><commit_msg>client test stuff<commit_after><|endoftext|>


<commit_before>package dirContent

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"path/filepath"
	"strings"
)

// Directory content struct
type Content struct {
	Name string
	Path string
	Size int64
	Type string
}

// Get directory Content and return
func getDirectoryInput(dir string) []byte {
	// Retrieve all the files in the input directory
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		log.Fatal(err)
	}

	var listOfContent []Content

	// Create a Content struct for each file and convert it to a JSON object.
	for _, file := range files {
		ext := filepath.Ext(file.Name())
		name := strings.TrimSuffix(file.Name(), ext)
		c := Content{
			Name: name,
			Path: dir,
			Size: file.Size(),
			Type: ext}

		listOfContent = append(listOfContent, c)
	}

	// Marshall a JSON-encoded string of the listOfContent
	mList, err := json.Marshal(listOfContent)
	if err != nil {
		log.Fatal(err)
	}

	return mList

}
<commit_msg>Function getting dir content and returns a JSON-encoded string renamed.<commit_after>package dirContent

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"path/filepath"
	"strings"
)

// Directory content struct
type Content struct {
	Name string
	Path string
	Size int64
	Type string
}

// Get directory Content and return
func getDirectoryContentInJSON(dir string) []byte {
	// Retrieve all the files in the input directory
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		log.Fatal(err)
	}

	var listOfContent []Content

	// Create a Content struct for each file and convert it to a JSON object.
	for _, file := range files {
		ext := filepath.Ext(file.Name())
		name := strings.TrimSuffix(file.Name(), ext)
		c := Content{
			Name: name,
			Path: dir,
			Size: file.Size(),
			Type: ext}

		listOfContent = append(listOfContent, c)
	}

	// Marshall a JSON-encoded string of the listOfContent
	mList, err := json.Marshal(listOfContent)
	if err != nil {
		log.Fatal(err)
	}

	return mList

}
<|endoftext|>


<commit_before>package filesystem

import (
	"github.com/abates/bms/database"
	"io"
	"os"
	"time"
)

type Asset interface {
	io.ReadWriteCloser
	io.Seeker
	ID() database.ID
	Mode() os.FileMode
	Name() string
	Owner() database.ID
	Readdir(count int) ([]os.FileInfo, error)
	SetName(string) error
	SetOwner(database.ID) error
	Stat() (os.FileInfo, error)
}

type AssetInfo interface {
	IsDir() bool
	Mode() os.FileMode
	ModTime() time.Time
	Name() string
	Size() int64
	Sys() interface{}
}
<commit_msg>No longer needed<commit_after><|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package filters

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	fctypesv1a1 "k8s.io/api/flowcontrol/v1alpha1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apiserver/pkg/apis/flowcontrol/bootstrap"
	"k8s.io/apiserver/pkg/authentication/user"
	apifilters "k8s.io/apiserver/pkg/endpoints/filters"
	epmetrics "k8s.io/apiserver/pkg/endpoints/metrics"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/server/mux"
	utilflowcontrol "k8s.io/apiserver/pkg/util/flowcontrol"
	fq "k8s.io/apiserver/pkg/util/flowcontrol/fairqueuing"
	fcmetrics "k8s.io/apiserver/pkg/util/flowcontrol/metrics"
	"k8s.io/component-base/metrics/legacyregistry"
)

type mockDecision int

const (
	decisionNoQueuingExecute mockDecision = iota
	decisionQueuingExecute
	decisionCancelWait
	decisionReject
	decisionSkipFilter
)

type fakeApfFilter struct {
	mockDecision mockDecision
	postEnqueue  func()
	postDequeue  func()
}

func (t fakeApfFilter) MaintainObservations(stopCh <-chan struct{}) {
}

func (t fakeApfFilter) Handle(ctx context.Context,
	requestDigest utilflowcontrol.RequestDigest,
	noteFn func(fs *fctypesv1a1.FlowSchema, pl *fctypesv1a1.PriorityLevelConfiguration),
	queueNoteFn fq.QueueNoteFn,
	execFn func(),
) {
	if t.mockDecision == decisionSkipFilter {
		panic("Handle should not be invoked")
	}
	noteFn(bootstrap.SuggestedFlowSchemaGlobalDefault, bootstrap.SuggestedPriorityLevelConfigurationGlobalDefault)
	switch t.mockDecision {
	case decisionNoQueuingExecute:
		execFn()
	case decisionQueuingExecute:
		queueNoteFn(true)
		t.postEnqueue()
		queueNoteFn(false)
		t.postDequeue()
		execFn()
	case decisionCancelWait:
		queueNoteFn(true)
		t.postEnqueue()
		queueNoteFn(false)
		t.postDequeue()
	case decisionReject:
		return
	}
}

func (t fakeApfFilter) Run(stopCh <-chan struct{}) error {
	return nil
}

func (t fakeApfFilter) Install(c *mux.PathRecorderMux) {
}

func newApfServerWithSingleRequest(decision mockDecision, t *testing.T) *httptest.Server {
	onExecuteFunc := func() {
		if decision == decisionCancelWait {
			t.Errorf("execute should not be invoked")
		}
		// atomicReadOnlyExecuting can be either 0 or 1 as we test one request at a time.
		if decision != decisionSkipFilter && atomicReadOnlyExecuting != 1 {
			t.Errorf("Wanted %d requests executing, got %d", 1, atomicReadOnlyExecuting)
		}
	}
	postExecuteFunc := func() {}
	// atomicReadOnlyWaiting can be either 0 or 1 as we test one request at a time.
	postEnqueueFunc := func() {
		if atomicReadOnlyWaiting != 1 {
			t.Errorf("Wanted %d requests in queue, got %d", 1, atomicReadOnlyWaiting)
		}
	}
	postDequeueFunc := func() {
		if atomicReadOnlyWaiting != 0 {
			t.Errorf("Wanted %d requests in queue, got %d", 0, atomicReadOnlyWaiting)
		}
	}
	return newApfServerWithHooks(decision, onExecuteFunc, postExecuteFunc, postEnqueueFunc, postDequeueFunc, t)
}

func newApfServerWithHooks(decision mockDecision, onExecute, postExecute, postEnqueue, postDequeue func(), t *testing.T) *httptest.Server {
	requestInfoFactory := &apirequest.RequestInfoFactory{APIPrefixes: sets.NewString("apis", "api"), GrouplessAPIPrefixes: sets.NewString("api")}
	longRunningRequestCheck := BasicLongRunningRequestCheck(sets.NewString("watch"), sets.NewString("proxy"))

	apfHandler := WithPriorityAndFairness(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		onExecute()
	}), longRunningRequestCheck, fakeApfFilter{
		mockDecision: decision,
		postEnqueue:  postEnqueue,
		postDequeue:  postDequeue,
	})

	handler := apifilters.WithRequestInfo(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r = r.WithContext(apirequest.WithUser(r.Context(), &user.DefaultInfo{
			Groups: []string{user.AllUnauthenticated},
		}))
		apfHandler.ServeHTTP(w, r)
		postExecute()
		if atomicReadOnlyExecuting != 0 {
			t.Errorf("Wanted %d requests executing, got %d", 0, atomicReadOnlyExecuting)
		}
	}), requestInfoFactory)

	apfServer := httptest.NewServer(handler)
	return apfServer
}

func TestApfSkipLongRunningRequest(t *testing.T) {
	epmetrics.Register()

	server := newApfServerWithSingleRequest(decisionSkipFilter, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	// send a watch request to test skipping long running request
	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces?watch=true", server.URL), http.StatusOK); err != nil {
		// request should not be rejected
		t.Error(err)
	}
}

func TestApfRejectRequest(t *testing.T) {
	epmetrics.Register()

	server := newApfServerWithSingleRequest(decisionReject, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusTooManyRequests); err != nil {
		t.Error(err)
	}

	checkForExpectedMetricsWithRetry(t, []string{
		"apiserver_request_terminations_total",
		"apiserver_dropped_requests_total",
	})
}

func TestApfExemptRequest(t *testing.T) {
	epmetrics.Register()
	fcmetrics.Register()

	// Wait for at least one sampling window to pass since creation of metrics.ReadWriteConcurrencyObserverPairGenerator,
	// so that an observation will cause some data to go into the Prometheus metrics.
	time.Sleep(time.Millisecond * 50)

	server := newApfServerWithSingleRequest(decisionNoQueuingExecute, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusOK); err != nil {
		t.Error(err)
	}

	checkForExpectedMetricsWithRetry(t, []string{
		"apiserver_current_inflight_requests",
		"apiserver_flowcontrol_read_vs_write_request_count_watermarks",
		"apiserver_flowcontrol_read_vs_write_request_count_samples",
	})
}

func TestApfExecuteRequest(t *testing.T) {
	epmetrics.Register()
	fcmetrics.Register()

	// Wait for at least one sampling window to pass since creation of metrics.ReadWriteConcurrencyObserverPairGenerator,
	// so that an observation will cause some data to go into the Prometheus metrics.
	time.Sleep(time.Millisecond * 50)

	server := newApfServerWithSingleRequest(decisionQueuingExecute, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusOK); err != nil {
		t.Error(err)
	}

	checkForExpectedMetricsWithRetry(t, []string{
		"apiserver_current_inflight_requests",
		"apiserver_current_inqueue_requests",
		"apiserver_flowcontrol_read_vs_write_request_count_watermarks",
		"apiserver_flowcontrol_read_vs_write_request_count_samples",
	})
}

func TestApfExecuteMultipleRequests(t *testing.T) {
	epmetrics.Register()
	fcmetrics.Register()

	// Wait for at least one sampling window to pass since creation of metrics.ReadWriteConcurrencyObserverPairGenerator,
	// so that an observation will cause some data to go into the Prometheus metrics.
	time.Sleep(time.Millisecond * 50)

	concurrentRequests := 5
	var preStartExecute, postStartExecute, preEnqueue, postEnqueue, preDequeue, postDequeue, finishExecute sync.WaitGroup
	for _, wg := range []*sync.WaitGroup{&preStartExecute, &postStartExecute, &preEnqueue, &postEnqueue, &preDequeue, &postDequeue, &finishExecute} {
		wg.Add(concurrentRequests)
	}

	onExecuteFunc := func() {
		preStartExecute.Done()
		preStartExecute.Wait()
		if int(atomicReadOnlyExecuting) != concurrentRequests {
			t.Errorf("Wanted %d requests executing, got %d", concurrentRequests, atomicReadOnlyExecuting)
		}
		postStartExecute.Done()
		postStartExecute.Wait()
	}

	postEnqueueFunc := func() {
		preEnqueue.Done()
		preEnqueue.Wait()
		if int(atomicReadOnlyWaiting) != concurrentRequests {
			t.Errorf("Wanted %d requests in queue, got %d", 1, atomicReadOnlyWaiting)

		}
		postEnqueue.Done()
		postEnqueue.Wait()
	}

	postDequeueFunc := func() {
		preDequeue.Done()
		preDequeue.Wait()
		if atomicReadOnlyWaiting != 0 {
			t.Errorf("Wanted %d requests in queue, got %d", 0, atomicReadOnlyWaiting)
		}
		postDequeue.Done()
		postDequeue.Wait()
	}

	postExecuteFunc := func() {
		finishExecute.Done()
		finishExecute.Wait()
	}

	server := newApfServerWithHooks(decisionQueuingExecute, onExecuteFunc, postExecuteFunc, postEnqueueFunc, postDequeueFunc, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	var wg sync.WaitGroup
	wg.Add(concurrentRequests)
	for i := 0; i < concurrentRequests; i++ {
		go func() {
			defer wg.Done()
			if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusOK); err != nil {
				t.Error(err)
			}
		}()
	}
	wg.Wait()

	checkForExpectedMetricsWithRetry(t, []string{
		"apiserver_current_inflight_requests",
		"apiserver_current_inqueue_requests",
		"apiserver_flowcontrol_read_vs_write_request_count_watermarks",
		"apiserver_flowcontrol_read_vs_write_request_count_samples",
	})
}

func TestApfCancelWaitRequest(t *testing.T) {
	epmetrics.Register()

	server := newApfServerWithSingleRequest(decisionCancelWait, t)
	defer server.Close()

	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusTooManyRequests); err != nil {
		t.Error(err)
	}

	checkForExpectedMetricsWithRetry(t, []string{
		"apiserver_current_inflight_requests",
		"apiserver_request_terminations_total",
		"apiserver_dropped_requests_total",
	})
}

// wait async metrics to be collected
func checkForExpectedMetricsWithRetry(t *testing.T, expectedMetrics []string) {
	maxRetries := 5
	var checkErrors []error
	for i := 0; i < maxRetries; i++ {
		t.Logf("Check for expected metrics with retry %d", i)
		metricsFamily, err := legacyregistry.DefaultGatherer.Gather()
		if err != nil {
			t.Fatalf("Failed to gather metrics %v", err)
		}

		metrics := map[string]interface{}{}
		for _, mf := range metricsFamily {
			metrics[*mf.Name] = mf
		}

		checkErrors = checkForExpectedMetrics(expectedMetrics, metrics)
		if checkErrors == nil {
			return
		}

		time.Sleep(1 * time.Second)
	}
	for _, checkError := range checkErrors {
		t.Error(checkError)
	}
}

func checkForExpectedMetrics(expectedMetrics []string, metrics map[string]interface{}) []error {
	var errs []error
	for _, metricName := range expectedMetrics {
		if _, ok := metrics[metricName]; !ok {
			if !ok {
				errs = append(errs, errors.New("Scraped metrics did not include expected metric "+metricName))
			}
		}
	}
	return errs
}
<commit_msg>removes filter metrics test retries<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package filters

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	fctypesv1a1 "k8s.io/api/flowcontrol/v1alpha1"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apiserver/pkg/apis/flowcontrol/bootstrap"
	"k8s.io/apiserver/pkg/authentication/user"
	apifilters "k8s.io/apiserver/pkg/endpoints/filters"
	epmetrics "k8s.io/apiserver/pkg/endpoints/metrics"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/server/mux"
	utilflowcontrol "k8s.io/apiserver/pkg/util/flowcontrol"
	fq "k8s.io/apiserver/pkg/util/flowcontrol/fairqueuing"
	fcmetrics "k8s.io/apiserver/pkg/util/flowcontrol/metrics"
	"k8s.io/component-base/metrics/legacyregistry"
)

type mockDecision int

const (
	decisionNoQueuingExecute mockDecision = iota
	decisionQueuingExecute
	decisionCancelWait
	decisionReject
	decisionSkipFilter
)

type fakeApfFilter struct {
	mockDecision mockDecision
	postEnqueue  func()
	postDequeue  func()
}

func (t fakeApfFilter) MaintainObservations(stopCh <-chan struct{}) {
}

func (t fakeApfFilter) Handle(ctx context.Context,
	requestDigest utilflowcontrol.RequestDigest,
	noteFn func(fs *fctypesv1a1.FlowSchema, pl *fctypesv1a1.PriorityLevelConfiguration),
	queueNoteFn fq.QueueNoteFn,
	execFn func(),
) {
	if t.mockDecision == decisionSkipFilter {
		panic("Handle should not be invoked")
	}
	noteFn(bootstrap.SuggestedFlowSchemaGlobalDefault, bootstrap.SuggestedPriorityLevelConfigurationGlobalDefault)
	switch t.mockDecision {
	case decisionNoQueuingExecute:
		execFn()
	case decisionQueuingExecute:
		queueNoteFn(true)
		t.postEnqueue()
		queueNoteFn(false)
		t.postDequeue()
		execFn()
	case decisionCancelWait:
		queueNoteFn(true)
		t.postEnqueue()
		queueNoteFn(false)
		t.postDequeue()
	case decisionReject:
		return
	}
}

func (t fakeApfFilter) Run(stopCh <-chan struct{}) error {
	return nil
}

func (t fakeApfFilter) Install(c *mux.PathRecorderMux) {
}

func newApfServerWithSingleRequest(decision mockDecision, t *testing.T) *httptest.Server {
	onExecuteFunc := func() {
		if decision == decisionCancelWait {
			t.Errorf("execute should not be invoked")
		}
		// atomicReadOnlyExecuting can be either 0 or 1 as we test one request at a time.
		if decision != decisionSkipFilter && atomicReadOnlyExecuting != 1 {
			t.Errorf("Wanted %d requests executing, got %d", 1, atomicReadOnlyExecuting)
		}
	}
	postExecuteFunc := func() {}
	// atomicReadOnlyWaiting can be either 0 or 1 as we test one request at a time.
	postEnqueueFunc := func() {
		if atomicReadOnlyWaiting != 1 {
			t.Errorf("Wanted %d requests in queue, got %d", 1, atomicReadOnlyWaiting)
		}
	}
	postDequeueFunc := func() {
		if atomicReadOnlyWaiting != 0 {
			t.Errorf("Wanted %d requests in queue, got %d", 0, atomicReadOnlyWaiting)
		}
	}
	return newApfServerWithHooks(decision, onExecuteFunc, postExecuteFunc, postEnqueueFunc, postDequeueFunc, t)
}

func newApfServerWithHooks(decision mockDecision, onExecute, postExecute, postEnqueue, postDequeue func(), t *testing.T) *httptest.Server {
	requestInfoFactory := &apirequest.RequestInfoFactory{APIPrefixes: sets.NewString("apis", "api"), GrouplessAPIPrefixes: sets.NewString("api")}
	longRunningRequestCheck := BasicLongRunningRequestCheck(sets.NewString("watch"), sets.NewString("proxy"))

	apfHandler := WithPriorityAndFairness(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		onExecute()
	}), longRunningRequestCheck, fakeApfFilter{
		mockDecision: decision,
		postEnqueue:  postEnqueue,
		postDequeue:  postDequeue,
	})

	handler := apifilters.WithRequestInfo(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r = r.WithContext(apirequest.WithUser(r.Context(), &user.DefaultInfo{
			Groups: []string{user.AllUnauthenticated},
		}))
		apfHandler.ServeHTTP(w, r)
		postExecute()
		if atomicReadOnlyExecuting != 0 {
			t.Errorf("Wanted %d requests executing, got %d", 0, atomicReadOnlyExecuting)
		}
	}), requestInfoFactory)

	apfServer := httptest.NewServer(handler)
	return apfServer
}

func TestApfSkipLongRunningRequest(t *testing.T) {
	epmetrics.Register()

	server := newApfServerWithSingleRequest(decisionSkipFilter, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	// send a watch request to test skipping long running request
	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces?watch=true", server.URL), http.StatusOK); err != nil {
		// request should not be rejected
		t.Error(err)
	}
}

func TestApfRejectRequest(t *testing.T) {
	epmetrics.Register()

	server := newApfServerWithSingleRequest(decisionReject, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusTooManyRequests); err != nil {
		t.Error(err)
	}

	checkForExpectedMetrics(t, []string{
		"apiserver_request_terminations_total",
		"apiserver_dropped_requests_total",
	})
}

func TestApfExemptRequest(t *testing.T) {
	epmetrics.Register()
	fcmetrics.Register()

	// Wait for at least one sampling window to pass since creation of metrics.ReadWriteConcurrencyObserverPairGenerator,
	// so that an observation will cause some data to go into the Prometheus metrics.
	time.Sleep(time.Millisecond * 50)

	server := newApfServerWithSingleRequest(decisionNoQueuingExecute, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusOK); err != nil {
		t.Error(err)
	}

	checkForExpectedMetrics(t, []string{
		"apiserver_current_inflight_requests",
		"apiserver_flowcontrol_read_vs_write_request_count_watermarks",
		"apiserver_flowcontrol_read_vs_write_request_count_samples",
	})
}

func TestApfExecuteRequest(t *testing.T) {
	epmetrics.Register()
	fcmetrics.Register()

	// Wait for at least one sampling window to pass since creation of metrics.ReadWriteConcurrencyObserverPairGenerator,
	// so that an observation will cause some data to go into the Prometheus metrics.
	time.Sleep(time.Millisecond * 50)

	server := newApfServerWithSingleRequest(decisionQueuingExecute, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusOK); err != nil {
		t.Error(err)
	}

	checkForExpectedMetrics(t, []string{
		"apiserver_current_inflight_requests",
		"apiserver_current_inqueue_requests",
		"apiserver_flowcontrol_read_vs_write_request_count_watermarks",
		"apiserver_flowcontrol_read_vs_write_request_count_samples",
	})
}

func TestApfExecuteMultipleRequests(t *testing.T) {
	epmetrics.Register()
	fcmetrics.Register()

	// Wait for at least one sampling window to pass since creation of metrics.ReadWriteConcurrencyObserverPairGenerator,
	// so that an observation will cause some data to go into the Prometheus metrics.
	time.Sleep(time.Millisecond * 50)

	concurrentRequests := 5
	preStartExecute, postStartExecute := &sync.WaitGroup{}, &sync.WaitGroup{}
	preEnqueue, postEnqueue := &sync.WaitGroup{}, &sync.WaitGroup{}
	preDequeue, postDequeue := &sync.WaitGroup{}, &sync.WaitGroup{}
	finishExecute := &sync.WaitGroup{}
	for _, wg := range []*sync.WaitGroup{preStartExecute, postStartExecute, preEnqueue, postEnqueue, preDequeue, postDequeue, finishExecute} {
		wg.Add(concurrentRequests)
	}

	onExecuteFunc := func() {
		preStartExecute.Done()
		preStartExecute.Wait()
		if int(atomicReadOnlyExecuting) != concurrentRequests {
			t.Errorf("Wanted %d requests executing, got %d", concurrentRequests, atomicReadOnlyExecuting)
		}
		postStartExecute.Done()
		postStartExecute.Wait()
	}

	postEnqueueFunc := func() {
		preEnqueue.Done()
		preEnqueue.Wait()
		if int(atomicReadOnlyWaiting) != concurrentRequests {
			t.Errorf("Wanted %d requests in queue, got %d", 1, atomicReadOnlyWaiting)

		}
		postEnqueue.Done()
		postEnqueue.Wait()
	}

	postDequeueFunc := func() {
		preDequeue.Done()
		preDequeue.Wait()
		if atomicReadOnlyWaiting != 0 {
			t.Errorf("Wanted %d requests in queue, got %d", 0, atomicReadOnlyWaiting)
		}
		postDequeue.Done()
		postDequeue.Wait()
	}

	postExecuteFunc := func() {
		finishExecute.Done()
		finishExecute.Wait()
	}

	server := newApfServerWithHooks(decisionQueuingExecute, onExecuteFunc, postExecuteFunc, postEnqueueFunc, postDequeueFunc, t)
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	StartPriorityAndFairnessWatermarkMaintenance(ctx.Done())

	var wg sync.WaitGroup
	wg.Add(concurrentRequests)
	for i := 0; i < concurrentRequests; i++ {
		go func() {
			defer wg.Done()
			if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusOK); err != nil {
				t.Error(err)
			}
		}()
	}
	wg.Wait()

	checkForExpectedMetrics(t, []string{
		"apiserver_current_inflight_requests",
		"apiserver_current_inqueue_requests",
		"apiserver_flowcontrol_read_vs_write_request_count_watermarks",
		"apiserver_flowcontrol_read_vs_write_request_count_samples",
	})
}

func TestApfCancelWaitRequest(t *testing.T) {
	epmetrics.Register()

	server := newApfServerWithSingleRequest(decisionCancelWait, t)
	defer server.Close()

	if err := expectHTTPGet(fmt.Sprintf("%s/api/v1/namespaces/default", server.URL), http.StatusTooManyRequests); err != nil {
		t.Error(err)
	}

	checkForExpectedMetrics(t, []string{
		"apiserver_current_inflight_requests",
		"apiserver_request_terminations_total",
		"apiserver_dropped_requests_total",
	})
}

// gathers and checks the metrics.
func checkForExpectedMetrics(t *testing.T, expectedMetrics []string) {
	metricsFamily, err := legacyregistry.DefaultGatherer.Gather()
	if err != nil {
		t.Fatalf("Failed to gather metrics %v", err)
	}

	metrics := map[string]interface{}{}
	for _, mf := range metricsFamily {
		metrics[*mf.Name] = mf
	}

	for _, metricName := range expectedMetrics {
		if _, ok := metrics[metricName]; !ok {
			if !ok {
				t.Errorf("Scraped metrics did not include expected metric %s", metricName)
			}
		}
	}
}
<|endoftext|>


<commit_before>package oneandone_cloudserver_api

import (
	"github.com/docker/machine/log"
	"net/http"
)

type FirewallPolicy struct {
	withId
	withName
	withDescription
	// should be fixed: Status        Status                   `json:"status"`
	Status        string                   `json:"state"`
	DefaultPolicy int                      `json:"default"`
	Rules         []FirewallPolicyRules    `json:"rules"`
	ServerIps     []FirewallPolicyServerIp `json:"server_ips"`
	withApi
}

type FirewallPolicyRules struct {
	withId
	Protocol string `json:"protocol"`
	PortFrom *int   `json:"port_from"`
	PortTo   *int   `json:"port_to"`
	SourceIp string `json:"source"`
}

type FirewallPolicyServerIp struct {
	withId
	Ip         string `json:"ip"`
	ServerName string `json:"server_name"`
}

type FirewallPolicyCreateData struct {
	Name        string                          `json:"name"`
	Description string                          `json:"description"`
	Rules       []FirewallPolicyRulesCreateData `json:"rules"`
}

type FirewallPolicyRulesCreateData struct {
	Protocol string `json:"protocol"`
	PortFrom *int   `json:"port_from"`
	PortTo   *int   `json:"port_to"`
	SourceIp string `json:"source"`
}

// GET /firewall_policies
func (api *API) GetFirewallPolicies() ([]FirewallPolicy, error) {
	log.Debug("requesting information about firewall policies")
	result := []FirewallPolicy{}
	err := api.Client.Get(createUrl(api, "firewall_policies"), &result, http.StatusOK)
	if err != nil {
		return nil, err
	}
	for index, _ := range result {
		result[index].api = api
	}
	return result, nil
}

// POST /firewall_policies
func (api *API) CreateFirewallPolicy(configuration FirewallPolicyCreateData) (*FirewallPolicy, error) {
	log.Debug("requesting to create a new firewall policy")
	result := new(FirewallPolicy)
	err := api.Client.Post(createUrl(api, "firewall_policies"), configuration, &result, http.StatusCreated)
	if err != nil {
		return nil, err
	}
	result.api = api
	return result, nil
}

// GET /firewall_policies/{id}
func (api *API) GetFirewallPolicy(Id string) (*FirewallPolicy, error) {
	log.Debug("requesting to about firewall policy ", Id)
	result := new(FirewallPolicy)
	err := api.Client.Get(createUrl(api, "firewall_policies", Id), &result, http.StatusOK)
	if err != nil {
		return nil, err
	}
	result.api = api
	return result, nil

}

// DELETE /firewall_policies/{id}
func (fwp *FirewallPolicy) Delete() (*FirewallPolicy, error) {
	log.Debug("Requested to delete firewall policy ", fwp.Id)
	result := new(FirewallPolicy)
	err := fwp.api.Client.Delete(createUrl(fwp.api, "firewall_policies", fwp.Id), &result, http.StatusOK)
	if err != nil {
		return nil, err
	}
	result.api = fwp.api
	return result, nil
}

// PUT /firewall_policies/{id}

// GET /firewall_policies/{id}/server_ips

// PUT /firewall_policies/{id}/server_ips

// GET /firewall_policies/{id}/server_ips/{id}

// DELETE /firewall_policies/{id}/server_ips/{id}

// GET /firewall_policies/{id}/rules

// PUT /firewall_policies/{id}/rules

// GET /firewall_policies/{id}/rules/{id}

// DELETE /firewall_policies/{id}/rules/{id}
<commit_msg>Add function to apply server IP to a firewall policy.<commit_after>package oneandone_cloudserver_api

import (
	"github.com/docker/machine/log"
	"net/http"
)

type FirewallPolicy struct {
	withId
	withName
	withDescription
	// should be fixed: Status        Status                   `json:"status"`
	Status        string                   `json:"state"`
	DefaultPolicy int                      `json:"default"`
	Rules         []FirewallPolicyRules    `json:"rules"`
	ServerIps     []FirewallPolicyServerIp `json:"server_ips"`
	withApi
}

type FirewallPolicyRules struct {
	withId
	Protocol string `json:"protocol"`
	PortFrom *int   `json:"port_from"`
	PortTo   *int   `json:"port_to"`
	SourceIp string `json:"source"`
}

type FirewallPolicyServerIp struct {
	withId
	Ip         string `json:"ip"`
	ServerName string `json:"server_name"`
}

type FirewallPolicyCreateData struct {
	Name        string                          `json:"name"`
	Description string                          `json:"description"`
	Rules       []FirewallPolicyRulesCreateData `json:"rules"`
}

type FirewallPolicyRulesCreateData struct {
	Protocol string `json:"protocol"`
	PortFrom *int   `json:"port_from"`
	PortTo   *int   `json:"port_to"`
	SourceIp string `json:"source"`
}

type FirewallPolicyAddIpsData struct {
	ServerIps []string `json:"server_ips"`
}

// GET /firewall_policies
func (api *API) GetFirewallPolicies() ([]FirewallPolicy, error) {
	log.Debug("requesting information about firewall policies")
	result := []FirewallPolicy{}
	err := api.Client.Get(createUrl(api, "firewall_policies"), &result, http.StatusOK)
	if err != nil {
		return nil, err
	}
	for index, _ := range result {
		result[index].api = api
	}
	return result, nil
}

// POST /firewall_policies
func (api *API) CreateFirewallPolicy(configuration FirewallPolicyCreateData) (*FirewallPolicy, error) {
	log.Debug("requesting to create a new firewall policy")
	result := new(FirewallPolicy)
	err := api.Client.Post(createUrl(api, "firewall_policies"), configuration, &result, http.StatusCreated)
	if err != nil {
		return nil, err
	}
	result.api = api
	return result, nil
}

// GET /firewall_policies/{id}
func (api *API) GetFirewallPolicy(Id string) (*FirewallPolicy, error) {
	log.Debug("requesting to about firewall policy ", Id)
	result := new(FirewallPolicy)
	err := api.Client.Get(createUrl(api, "firewall_policies", Id), &result, http.StatusOK)
	if err != nil {
		return nil, err
	}
	result.api = api
	return result, nil

}

// DELETE /firewall_policies/{id}
func (fwp *FirewallPolicy) Delete() (*FirewallPolicy, error) {
	log.Debug("Requested to delete firewall policy ", fwp.Id)
	result := new(FirewallPolicy)
	err := fwp.api.Client.Delete(createUrl(fwp.api, "firewall_policies", fwp.Id), &result, http.StatusOK)
	if err != nil {
		return nil, err
	}
	result.api = fwp.api
	return result, nil
}

// PUT /firewall_policies/{id}

// GET /firewall_policies/{id}/server_ips

// PUT /firewall_policies/{id}/server_ips
func (fwp *FirewallPolicy) AddIp(ipId string) (*FirewallPolicy, error) {
	log.Debugf("Requested to apply firewall policy '%v' to ip '%v'", fwp.Id, ipId)
	result := new(FirewallPolicy)
	request := FirewallPolicyAddIpsData{
		ServerIps: []string{ipId},
	}
	err := fwp.api.Client.Put(createUrl(fwp.api, "firewall_policies", fwp.Id, "server_ips"), request, result, http.StatusAccepted)
	if err != nil {
		return nil, err
	}
	result.api = fwp.api
	return result, nil
}

// GET /firewall_policies/{id}/server_ips/{id}

// DELETE /firewall_policies/{id}/server_ips/{id}

// GET /firewall_policies/{id}/rules

// PUT /firewall_policies/{id}/rules

// GET /firewall_policies/{id}/rules/{id}

// DELETE /firewall_policies/{id}/rules/{id}
<|endoftext|>


<commit_before>// Package pastebin is a simple modern and powerful pastebin service
package main

import (
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"html"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	// uniuri is used for easy random string generation
	"github.com/dchest/uniuri"
	// pygments is used for syntax highlighting
	"github.com/ewhal/pygments"
	// mysql driver
	_ "github.com/go-sql-driver/mysql"
	// mux is used for url routing
	"github.com/gorilla/mux"
)

const (
	// ADDRESS that pastebin will return links for
	ADDRESS = "http://localhost:9900"
	// LENGTH of paste id
	LENGTH = 6
	// PORT that pastebin will listen on
	PORT = ":9900"
	// USERNAME for database
	USERNAME = ""
	// PASS database password
	PASS = ""
	// NAME database name
	NAME = ""
	// DATABASE connection String
	DATABASE = USERNAME + ":" + PASS + "@/" + NAME + "?charset=utf8"
)

// Template pages
var templates = template.Must(template.ParseFiles("assets/paste.html", "assets/index.html", "assets/clone.html"))
var syntax, _ = ioutil.ReadFile("assets/syntax.html")

// Response API struct
type Response struct {
	ID     string `json:"id"`
	TITLE  string `json:"title"`
	HASH   string `json:"hash"`
	URL    string `json:"url"`
	SIZE   int    `json:"size"`
	DELKEY string `json:"delkey"`
}

// Page generation struct
type Page struct {
	Title    string
	Body     []byte
	Raw      string
	Home     string
	Download string
	Clone    string
}

// check error handling function
func check(err error) {
	if err != nil {
		fmt.Println(err)
	}
}

// generateName uses uniuri to generate a random string that isn't in the
// database
func generateName() string {
	// use uniuri to generate random string
	id := uniuri.NewLen(LENGTH)

	db, err := sql.Open("mysql", DATABASE)
	check(err)
	defer db.Close()
	// query database if id exists and if it does call generateName again
	query, err := db.Query("select id from pastebin where id=?", id)
	if err != sql.ErrNoRows {
		for query.Next() {
			generateName()
		}
	}

	return id

}

// hash hashes paste into a sha1 hash
func hash(paste string) string {
	hasher := sha1.New()

	hasher.Write([]byte(paste))
	sha := base64.URLEncoding.EncodeToString(hasher.Sum(nil))
	return sha
}

// durationFromExpiry takes the expiry in string format and returns the duration
// that the paste will exist for
func durationFromExpiry(expiry string) time.Duration {
	switch expiry {
	case "5 minutes":
		return time.Minute * 5
	case "1 hour":
		return time.Hour + 1 // XXX: did you mean '*'?
	case "1 day":
		return time.Hour * 24
	case "1 week":
		return time.Hour * 24 * 7
	case "1 month":
		return time.Hour * 24 * 30
	case "1 year":
		return time.Hour * 24 * 365
	case "forever":
		return time.Hour * 24 * (365 * 20)
	}
	return time.Hour * 24 * (365 * 20)
}

// save function handles the saving of each paste.
// raw string is the raw paste input
// lang string is the user specified language for syntax highlighting
// title string user customized title
// expiry string duration that the paste will exist for
// Returns Response struct
func save(raw string, lang string, title string, expiry string) Response {

	db, err := sql.Open("mysql", DATABASE)
	check(err)
	defer db.Close()

	// hash paste data and query database to see if paste exists
	sha := hash(raw)
	query, err := db.Query("select id, title, hash, data, delkey from pastebin where hash=?", sha)

	if err != sql.ErrNoRows {
		for query.Next() {
			var id, title, hash, paste, delkey string
			err := query.Scan(&id, &title, &hash, &paste, &delkey)
			check(err)
			url := ADDRESS + "/p/" + id
			return Response{id, title, hash, url, len(paste), delkey}
		}
	}
	id := generateName()
	url := ADDRESS + "/p/" + id
	if lang != "" {
		url += "/" + lang
	}

	const timeFormat = "2006-01-02 15:04:05"
	expiryTime := time.Now().Add(durationFromExpiry(expiry)).Format(timeFormat)

	delKey := uniuri.NewLen(40)
	dataEscaped := html.EscapeString(raw)

	stmt, err := db.Prepare("INSERT INTO pastebin(id, title, hash, data, delkey, expiry) values(?,?,?,?,?,?)")
	check(err)
	if title == "" {
		title = id
	}
	_, err = stmt.Exec(id, html.EscapeString(title), sha, dataEscaped, delKey, expiryTime)
	check(err)

	return Response{id, title, sha, url, len(dataEscaped), delKey}
}

// delHandler checks to see if delkey and pasteid exist in the database.
// if both exist and are correct the paste will be removed.
func delHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["pasteId"]
	delkey := vars["delKey"]

	db, err := sql.Open("mysql", DATABASE)
	check(err)
	defer db.Close()

	stmt, err := db.Prepare("delete from pastebin where delkey=? and id=?")
	check(err)

	res, err := stmt.Exec(html.EscapeString(delkey), html.EscapeString(id))
	check(err)

	_, err = res.RowsAffected()
	if err != sql.ErrNoRows {
		io.WriteString(w, id+" deleted")
	}
}

// saveHandler
func saveHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	output := vars["output"]
	switch r.Method {
	case "POST":
		paste := r.FormValue("p")
		lang := r.FormValue("lang")
		title := r.FormValue("title")
		expiry := r.FormValue("expiry")
		if paste == "" {
			http.Error(w, "Empty paste", 500)
			return
		}
		b := save(paste, lang, title, expiry)

		switch output {
		case "json":
			w.Header().Set("Content-Type", "application/json")
			err := json.NewEncoder(w).Encode(b)
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

		case "xml":
			x, err := xml.MarshalIndent(b, "", "  ")
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

			w.Header().Set("Content-Type", "application/xml")
			w.Write(x)

		case "html":
			w.Header().Set("Content-Type", "text/html")
			io.WriteString(w, "
URL: "+b.URL+"
")
			io.WriteString(w, "
Delete Key: "+b.DELKEY+"
")

		case "redirect":
			http.Redirect(w, r, b.URL, 301)

		default:
			w.Header().Set("Content-Type", "plain/text")
			io.WriteString(w, b.URL+"\n")
			io.WriteString(w, "delete key: "+b.DELKEY+"\n")
		}
	}

}

// highlight uses user specified input to call pygments library to highlight the
// paste
func highlight(s string, lang string) (string, error) {

	highlight, err := pygments.Highlight(html.UnescapeString(s), html.EscapeString(lang), "html", "style=autumn,linenos=True, lineanchors=True,anchorlinenos=True,noclasses=True,", "utf-8")
	if err != nil {
		return "", err
	}
	return highlight, nil

}

// getPaste takes pasteid and language
// queries the database and returns paste data
func getPaste(paste string, lang string) (string, string) {
	param1 := html.EscapeString(paste)
	db, err := sql.Open("mysql", DATABASE)
	check(err)
	defer db.Close()
	var title, s string
	var expiry string
	err = db.QueryRow("select title, data, expiry from pastebin where id=?", param1).Scan(&title, &s, &expiry)
	check(err)
	if time.Now().Format("2006-01-02 15:04:05") > expiry {
		stmt, err := db.Prepare("delete from pastebin where id=?")
		check(err)
		_, err = stmt.Exec(param1)
		check(err)
		return "Error invalid paste", ""
	}

	if err == sql.ErrNoRows {
		return "Error invalid paste", ""
	}
	if lang != "" {
		high, err := highlight(s, lang)
		check(err)
		return high, html.UnescapeString(title)
	}
	return html.UnescapeString(s), html.UnescapeString(title)
}

func pasteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	lang := vars["lang"]

	s, title := getPaste(paste, lang)

	// button links
	link := ADDRESS + "/raw/" + paste
	download := ADDRESS + "/download/" + paste
	clone := ADDRESS + "/clone/" + paste
	// Page struct
	p := &Page{
		Title:    title,
		Body:     []byte(s),
		Raw:      link,
		Home:     ADDRESS,
		Download: download,
		Clone:    clone,
	}
	if lang == "" {

		err := templates.ExecuteTemplate(w, "paste.html", p)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}

	} else {
		fmt.Fprintf(w, string(syntax), p.Title, p.Title, s, p.Home, p.Download, p.Raw, p.Clone)

	}
}

func cloneHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]

	s, title := getPaste(paste, "")

	// Page links
	link := ADDRESS + "/raw/" + paste
	download := ADDRESS + "/download/" + paste
	clone := ADDRESS + "/clone/" + paste

	// Clone page struct
	p := &Page{
		Title:    title,
		Body:     []byte(s),
		Raw:      link,
		Home:     ADDRESS,
		Download: download,
		Clone:    clone,
	}
	err := templates.ExecuteTemplate(w, "clone.html", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

}
func downloadHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, _ := getPaste(paste, "")

	// Set header to an attachment so browser will automatically download it
	w.Header().Set("Content-Disposition", "attachment; filename="+paste)
	w.Header().Set("Content-Type", r.Header.Get("Content-Type"))
	io.WriteString(w, s)

}
func rawHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, _ := getPaste(paste, "")

	w.Header().Set("Content-Type", "plain/text")
	// simply write string to browser
	io.WriteString(w, s)

}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	err := templates.ExecuteTemplate(w, "index.html", &Page{})
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func main() {
	router := mux.NewRouter()
	router.HandleFunc("/p/{pasteId}", pasteHandler)
	router.HandleFunc("/raw/{pasteId}", rawHandler)
	router.HandleFunc("/p/{pasteId}/{lang}", pasteHandler)
	router.HandleFunc("/clone/{pasteId}", cloneHandler)
	router.HandleFunc("/download/{pasteId}", downloadHandler)
	router.HandleFunc("/save", saveHandler)
	router.HandleFunc("/save/{output}", saveHandler)
	router.HandleFunc("/del/{pasteId}/{delKey}", delHandler)
	router.HandleFunc("/", rootHandler)
	err := http.ListenAndServe(PORT, router)
	if err != nil {
		log.Fatal(err)
	}

}
<commit_msg>Force raw to be text/plain<commit_after>// Package pastebin is a simple modern and powerful pastebin service
package main

import (
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"html"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	// uniuri is used for easy random string generation
	"github.com/dchest/uniuri"
	// pygments is used for syntax highlighting
	"github.com/ewhal/pygments"
	// mysql driver
	_ "github.com/go-sql-driver/mysql"
	// mux is used for url routing
	"github.com/gorilla/mux"
)

const (
	// ADDRESS that pastebin will return links for
	ADDRESS = "http://localhost:9900"
	// LENGTH of paste id
	LENGTH = 6
	// PORT that pastebin will listen on
	PORT = ":9900"
	// USERNAME for database
	USERNAME = ""
	// PASS database password
	PASS = ""
	// NAME database name
	NAME = ""
	// DATABASE connection String
	DATABASE = USERNAME + ":" + PASS + "@/" + NAME + "?charset=utf8"
)

// Template pages
var templates = template.Must(template.ParseFiles("assets/paste.html", "assets/index.html", "assets/clone.html"))
var syntax, _ = ioutil.ReadFile("assets/syntax.html")

// Response API struct
type Response struct {
	ID     string `json:"id"`
	TITLE  string `json:"title"`
	HASH   string `json:"hash"`
	URL    string `json:"url"`
	SIZE   int    `json:"size"`
	DELKEY string `json:"delkey"`
}

// Page generation struct
type Page struct {
	Title    string
	Body     []byte
	Raw      string
	Home     string
	Download string
	Clone    string
}

// check error handling function
func check(err error) {
	if err != nil {
		fmt.Println(err)
	}
}

// generateName uses uniuri to generate a random string that isn't in the
// database
func generateName() string {
	// use uniuri to generate random string
	id := uniuri.NewLen(LENGTH)

	db, err := sql.Open("mysql", DATABASE)
	check(err)
	defer db.Close()
	// query database if id exists and if it does call generateName again
	query, err := db.Query("select id from pastebin where id=?", id)
	if err != sql.ErrNoRows {
		for query.Next() {
			generateName()
		}
	}

	return id

}

// hash hashes paste into a sha1 hash
func hash(paste string) string {
	hasher := sha1.New()

	hasher.Write([]byte(paste))
	sha := base64.URLEncoding.EncodeToString(hasher.Sum(nil))
	return sha
}

// durationFromExpiry takes the expiry in string format and returns the duration
// that the paste will exist for
func durationFromExpiry(expiry string) time.Duration {
	switch expiry {
	case "5 minutes":
		return time.Minute * 5
	case "1 hour":
		return time.Hour + 1 // XXX: did you mean '*'?
	case "1 day":
		return time.Hour * 24
	case "1 week":
		return time.Hour * 24 * 7
	case "1 month":
		return time.Hour * 24 * 30
	case "1 year":
		return time.Hour * 24 * 365
	case "forever":
		return time.Hour * 24 * (365 * 20)
	}
	return time.Hour * 24 * (365 * 20)
}

// save function handles the saving of each paste.
// raw string is the raw paste input
// lang string is the user specified language for syntax highlighting
// title string user customized title
// expiry string duration that the paste will exist for
// Returns Response struct
func save(raw string, lang string, title string, expiry string) Response {

	db, err := sql.Open("mysql", DATABASE)
	check(err)
	defer db.Close()

	// hash paste data and query database to see if paste exists
	sha := hash(raw)
	query, err := db.Query("select id, title, hash, data, delkey from pastebin where hash=?", sha)

	if err != sql.ErrNoRows {
		for query.Next() {
			var id, title, hash, paste, delkey string
			err := query.Scan(&id, &title, &hash, &paste, &delkey)
			check(err)
			url := ADDRESS + "/p/" + id
			return Response{id, title, hash, url, len(paste), delkey}
		}
	}
	id := generateName()
	url := ADDRESS + "/p/" + id
	if lang != "" {
		url += "/" + lang
	}

	const timeFormat = "2006-01-02 15:04:05"
	expiryTime := time.Now().Add(durationFromExpiry(expiry)).Format(timeFormat)

	delKey := uniuri.NewLen(40)
	dataEscaped := html.EscapeString(raw)

	stmt, err := db.Prepare("INSERT INTO pastebin(id, title, hash, data, delkey, expiry) values(?,?,?,?,?,?)")
	check(err)
	if title == "" {
		title = id
	}
	_, err = stmt.Exec(id, html.EscapeString(title), sha, dataEscaped, delKey, expiryTime)
	check(err)

	return Response{id, title, sha, url, len(dataEscaped), delKey}
}

// delHandler checks to see if delkey and pasteid exist in the database.
// if both exist and are correct the paste will be removed.
func delHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	id := vars["pasteId"]
	delkey := vars["delKey"]

	db, err := sql.Open("mysql", DATABASE)
	check(err)
	defer db.Close()

	stmt, err := db.Prepare("delete from pastebin where delkey=? and id=?")
	check(err)

	res, err := stmt.Exec(html.EscapeString(delkey), html.EscapeString(id))
	check(err)

	_, err = res.RowsAffected()
	if err != sql.ErrNoRows {
		io.WriteString(w, id+" deleted")
	}
}

// saveHandler
func saveHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	output := vars["output"]
	switch r.Method {
	case "POST":
		paste := r.FormValue("p")
		lang := r.FormValue("lang")
		title := r.FormValue("title")
		expiry := r.FormValue("expiry")
		if paste == "" {
			http.Error(w, "Empty paste", 500)
			return
		}
		b := save(paste, lang, title, expiry)

		switch output {
		case "json":
			w.Header().Set("Content-Type", "application/json")
			err := json.NewEncoder(w).Encode(b)
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

		case "xml":
			x, err := xml.MarshalIndent(b, "", "  ")
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

			w.Header().Set("Content-Type", "application/xml")
			w.Write(x)

		case "html":
			w.Header().Set("Content-Type", "text/html")
			io.WriteString(w, "
URL: "+b.URL+"
")
			io.WriteString(w, "
Delete Key: "+b.DELKEY+"
")

		case "redirect":
			http.Redirect(w, r, b.URL, 301)

		default:
			w.Header().Set("Content-Type", "plain/text")
			io.WriteString(w, b.URL+"\n")
			io.WriteString(w, "delete key: "+b.DELKEY+"\n")
		}
	}

}

// highlight uses user specified input to call pygments library to highlight the
// paste
func highlight(s string, lang string) (string, error) {

	highlight, err := pygments.Highlight(html.UnescapeString(s), html.EscapeString(lang), "html", "style=autumn,linenos=True, lineanchors=True,anchorlinenos=True,noclasses=True,", "utf-8")
	if err != nil {
		return "", err
	}
	return highlight, nil

}

// getPaste takes pasteid and language
// queries the database and returns paste data
func getPaste(paste string, lang string) (string, string) {
	param1 := html.EscapeString(paste)
	db, err := sql.Open("mysql", DATABASE)
	check(err)
	defer db.Close()
	var title, s string
	var expiry string
	err = db.QueryRow("select title, data, expiry from pastebin where id=?", param1).Scan(&title, &s, &expiry)
	check(err)
	if time.Now().Format("2006-01-02 15:04:05") > expiry {
		stmt, err := db.Prepare("delete from pastebin where id=?")
		check(err)
		_, err = stmt.Exec(param1)
		check(err)
		return "Error invalid paste", ""
	}

	if err == sql.ErrNoRows {
		return "Error invalid paste", ""
	}
	if lang != "" {
		high, err := highlight(s, lang)
		check(err)
		return high, html.UnescapeString(title)
	}
	return html.UnescapeString(s), html.UnescapeString(title)
}

func pasteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	lang := vars["lang"]

	s, title := getPaste(paste, lang)

	// button links
	link := ADDRESS + "/raw/" + paste
	download := ADDRESS + "/download/" + paste
	clone := ADDRESS + "/clone/" + paste
	// Page struct
	p := &Page{
		Title:    title,
		Body:     []byte(s),
		Raw:      link,
		Home:     ADDRESS,
		Download: download,
		Clone:    clone,
	}
	if lang == "" {

		err := templates.ExecuteTemplate(w, "paste.html", p)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}

	} else {
		fmt.Fprintf(w, string(syntax), p.Title, p.Title, s, p.Home, p.Download, p.Raw, p.Clone)

	}
}

func cloneHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]

	s, title := getPaste(paste, "")

	// Page links
	link := ADDRESS + "/raw/" + paste
	download := ADDRESS + "/download/" + paste
	clone := ADDRESS + "/clone/" + paste

	// Clone page struct
	p := &Page{
		Title:    title,
		Body:     []byte(s),
		Raw:      link,
		Home:     ADDRESS,
		Download: download,
		Clone:    clone,
	}
	err := templates.ExecuteTemplate(w, "clone.html", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

}
func downloadHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, _ := getPaste(paste, "")

	// Set header to an attachment so browser will automatically download it
	w.Header().Set("Content-Disposition", "attachment; filename="+paste)
	w.Header().Set("Content-Type", r.Header.Get("Content-Type"))
	io.WriteString(w, s)

}
func rawHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, _ := getPaste(paste, "")

	w.Header().Set("Content-Type", "text/plain; charset=UTF-8; imeanit=yes")
	// simply write string to browser
	io.WriteString(w, s)

}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	err := templates.ExecuteTemplate(w, "index.html", &Page{})
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func main() {
	router := mux.NewRouter()
	router.HandleFunc("/p/{pasteId}", pasteHandler)
	router.HandleFunc("/raw/{pasteId}", rawHandler)
	router.HandleFunc("/p/{pasteId}/{lang}", pasteHandler)
	router.HandleFunc("/clone/{pasteId}", cloneHandler)
	router.HandleFunc("/download/{pasteId}", downloadHandler)
	router.HandleFunc("/save", saveHandler)
	router.HandleFunc("/save/{output}", saveHandler)
	router.HandleFunc("/del/{pasteId}/{delKey}", delHandler)
	router.HandleFunc("/", rootHandler)
	err := http.ListenAndServe(PORT, router)
	if err != nil {
		log.Fatal(err)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"github.com/miekg/dns"
	"os"
	"strings"
	"time"
)

const (
	TIMEOUT time.Duration = 5 // seconds
)

var (
	localm *dns.Msg
	localc *dns.Client
	conf   *dns.ClientConfig
)

func localQuery(qname string, qtype uint16) (r *dns.Msg, err error) {
	localm.SetQuestion(qname, qtype)
	for i := range conf.Servers {
		server := conf.Servers[i]
		r, err := localc.Exchange(localm, server+":"+conf.Port)
		if r == nil || r.Rcode == dns.RcodeNameError || r.Rcode == dns.RcodeSuccess {
			return r, err
		}
	}
	return nil, errors.New("No name server to answer the question")
}

func main() {
	var err error
	if len(os.Args) != 2 {
		fmt.Printf("%s ZONE\n", os.Args[0])
		os.Exit(1)
	}
	conf, err = dns.ClientConfigFromFile("/etc/resolv.conf")
	if conf == nil {
		fmt.Printf("Cannot initialize the local resolver: %s\n", err)
		os.Exit(1)
	}
	localm = new(dns.Msg)
	localm.MsgHdr.RecursionDesired = true
	localm.Question = make([]dns.Question, 1)
	localc = new(dns.Client)
	localc.ReadTimeout = TIMEOUT * 1e9
	r, err := localQuery(dns.Fqdn(os.Args[1]), dns.TypeNS)
	if r == nil {
		fmt.Printf("Cannot retrieve the list of name servers for %s: %s\n", dns.Fqdn(os.Args[1]), err)
		os.Exit(1)
	}
	if r.Rcode == dns.RcodeNameError {
		fmt.Printf("No such domain %s\n", dns.Fqdn(os.Args[1]))
		os.Exit(1)
	}
	m := new(dns.Msg)
	m.MsgHdr.RecursionDesired = false
	m.Question = make([]dns.Question, 1)
	c := new(dns.Client)
	c.ReadTimeout = TIMEOUT * 1e9
	success := true
	numNS := 0
	for _, ans := range r.Answer {
		switch ans.(type) {
		case *dns.RR_NS:
			nameserver := ans.(*dns.RR_NS).Ns
			numNS += 1
			ips := make([]string, 0)
			fmt.Printf("%s : ", nameserver)
			ra, err := localQuery(nameserver, dns.TypeA)
			if ra == nil {
				fmt.Printf("Error getting the IPv4 address of %s: %s\n", nameserver, err)
				os.Exit(1)
			}
			if ra.Rcode != dns.RcodeSuccess {
				fmt.Printf("Error getting the IPv4 address of %s: %s\n", nameserver, dns.Rcode_str[ra.Rcode])
				os.Exit(1)
			}
			for _, ansa := range ra.Answer {
				switch ansa.(type) {
				case *dns.RR_A:
					ips = append(ips, ansa.(*dns.RR_A).A.String())
				}
			}
			raaaa, err := localQuery(nameserver, dns.TypeAAAA)
			if raaaa == nil {
				fmt.Printf("Error getting the IPv6 address of %s: %s\n", nameserver, err)
				os.Exit(1)
			}
			if raaaa.Rcode != dns.RcodeSuccess {
				fmt.Printf("Error getting the IPv6 address of %s: %s\n", nameserver, dns.Rcode_str[raaaa.Rcode])
				os.Exit(1)
			}
			for _, ansaaaa := range raaaa.Answer {
				switch ansaaaa.(type) {
				case *dns.RR_AAAA:
					ips = append(ips, ansaaaa.(*dns.RR_AAAA).AAAA.String())
				}
			}
			if len(ips) == 0 {
				success = false
				fmt.Printf("No IP address for this server")
			}
			for _, ip := range ips {
				m.Question[0] = dns.Question{dns.Fqdn(os.Args[1]), dns.TypeSOA, dns.ClassINET}
				nsAddressPort := ""
				if strings.ContainsAny(":", ip) {
					/* IPv6 address */
					nsAddressPort = "[" + ip + "]:53"
				} else {
					nsAddressPort = ip + ":53"
				}
				soa, err := c.Exchange(m, nsAddressPort)
				// TODO: retry if timeout? Otherwise, one lost UDP packet and it is the end
				if soa == nil {
					success = false
					fmt.Printf("%s (%s) ", ip, err)
				} else {
					if soa.Rcode != dns.RcodeSuccess {
						success = false
						fmt.Printf("%s (%s) ", ips, dns.Rcode_str[soa.Rcode])
					} else {
						if len(soa.Answer) == 0 { /* May happen if the server is a recursor, not authoritative, since we query with RD=0 */
							success = false
							fmt.Printf("%s (0 answer) ", ip)
						} else {
							rsoa := soa.Answer[0]
							switch rsoa.(type) {
							case *dns.RR_SOA:
								if soa.MsgHdr.Authoritative {
									/* TODO: test if all name servers have the same serial ? */
									fmt.Printf("%s (%d) ", ips, rsoa.(*dns.RR_SOA).Serial)
								} else {
									success = false
									fmt.Printf("%s (not authoritative) ", ips)
								}
							}
						}
					}
				}
			}
			fmt.Printf("\n")
		}
	}
	if numNS == 0 {
		fmt.Printf("No NS records for \"%s\". It is probably a CNAME to a domain but not a zone\n", dns.Fqdn(os.Args[1]))
		os.Exit(1)
	}
	if success {
		os.Exit(0)
	}
	os.Exit(1)
}
<commit_msg>Go comments<commit_after>package main

import (
	"errors"
	"fmt"
	"github.com/miekg/dns"
	"os"
	"strings"
	"time"
)

const (
	TIMEOUT time.Duration = 5 // seconds
)

var (
	localm *dns.Msg
	localc *dns.Client
	conf   *dns.ClientConfig
)

func localQuery(qname string, qtype uint16) (r *dns.Msg, err error) {
	localm.SetQuestion(qname, qtype)
	for i := range conf.Servers {
		server := conf.Servers[i]
		r, err := localc.Exchange(localm, server+":"+conf.Port)
		if r == nil || r.Rcode == dns.RcodeNameError || r.Rcode == dns.RcodeSuccess {
			return r, err
		}
	}
	return nil, errors.New("No name server to answer the question")
}

func main() {
	var err error
	if len(os.Args) != 2 {
		fmt.Printf("%s ZONE\n", os.Args[0])
		os.Exit(1)
	}
	conf, err = dns.ClientConfigFromFile("/etc/resolv.conf")
	if conf == nil {
		fmt.Printf("Cannot initialize the local resolver: %s\n", err)
		os.Exit(1)
	}
	localm = new(dns.Msg)
	localm.MsgHdr.RecursionDesired = true
	localm.Question = make([]dns.Question, 1)
	localc = new(dns.Client)
	localc.ReadTimeout = TIMEOUT * 1e9
	r, err := localQuery(dns.Fqdn(os.Args[1]), dns.TypeNS)
	if r == nil {
		fmt.Printf("Cannot retrieve the list of name servers for %s: %s\n", dns.Fqdn(os.Args[1]), err)
		os.Exit(1)
	}
	if r.Rcode == dns.RcodeNameError {
		fmt.Printf("No such domain %s\n", dns.Fqdn(os.Args[1]))
		os.Exit(1)
	}
	m := new(dns.Msg)
	m.MsgHdr.RecursionDesired = false
	m.Question = make([]dns.Question, 1)
	c := new(dns.Client)
	c.ReadTimeout = TIMEOUT * 1e9
	success := true
	numNS := 0
	for _, ans := range r.Answer {
		switch ans.(type) {
		case *dns.RR_NS:
			nameserver := ans.(*dns.RR_NS).Ns
			numNS += 1
			ips := make([]string, 0)
			fmt.Printf("%s : ", nameserver)
			ra, err := localQuery(nameserver, dns.TypeA)
			if ra == nil {
				fmt.Printf("Error getting the IPv4 address of %s: %s\n", nameserver, err)
				os.Exit(1)
			}
			if ra.Rcode != dns.RcodeSuccess {
				fmt.Printf("Error getting the IPv4 address of %s: %s\n", nameserver, dns.Rcode_str[ra.Rcode])
				os.Exit(1)
			}
			for _, ansa := range ra.Answer {
				switch ansa.(type) {
				case *dns.RR_A:
					ips = append(ips, ansa.(*dns.RR_A).A.String())
				}
			}
			raaaa, err := localQuery(nameserver, dns.TypeAAAA)
			if raaaa == nil {
				fmt.Printf("Error getting the IPv6 address of %s: %s\n", nameserver, err)
				os.Exit(1)
			}
			if raaaa.Rcode != dns.RcodeSuccess {
				fmt.Printf("Error getting the IPv6 address of %s: %s\n", nameserver, dns.Rcode_str[raaaa.Rcode])
				os.Exit(1)
			}
			for _, ansaaaa := range raaaa.Answer {
				switch ansaaaa.(type) {
				case *dns.RR_AAAA:
					ips = append(ips, ansaaaa.(*dns.RR_AAAA).AAAA.String())
				}
			}
			if len(ips) == 0 {
				success = false
				fmt.Printf("No IP address for this server")
			}
			for _, ip := range ips {
				m.Question[0] = dns.Question{dns.Fqdn(os.Args[1]), dns.TypeSOA, dns.ClassINET}
				nsAddressPort := ""
				if strings.ContainsAny(":", ip) {
					/* IPv6 address */
					nsAddressPort = "[" + ip + "]:53"
				} else {
					nsAddressPort = ip + ":53"
				}
				soa, err := c.Exchange(m, nsAddressPort)
				// TODO: retry if timeout? Otherwise, one lost UDP packet and it is the end
				if soa == nil {
					success = false
					fmt.Printf("%s (%s) ", ip, err)
				} else {
					if soa.Rcode != dns.RcodeSuccess {
						success = false
						fmt.Printf("%s (%s) ", ips, dns.Rcode_str[soa.Rcode])
					} else {
						if len(soa.Answer) == 0 { // May happen if the server is a recursor, not authoritative, since we query with RD=0
							success = false
							fmt.Printf("%s (0 answer) ", ip)
						} else {
							rsoa := soa.Answer[0]
							switch rsoa.(type) {
							case *dns.RR_SOA:
								if soa.MsgHdr.Authoritative {
									// TODO: test if all name servers have the same serial ?
									fmt.Printf("%s (%d) ", ips, rsoa.(*dns.RR_SOA).Serial)
								} else {
									success = false
									fmt.Printf("%s (not authoritative) ", ips)
								}
							}
						}
					}
				}
			}
			fmt.Printf("\n")
		}
	}
	if numNS == 0 {
		fmt.Printf("No NS records for \"%s\". It is probably a CNAME to a domain but not a zone\n", dns.Fqdn(os.Args[1]))
		os.Exit(1)
	}
	if success {
		os.Exit(0)
	}
	os.Exit(1)
}
<|endoftext|>