<commit_before>package main

import (
	"fmt"

	"github.com/whitepages/terraform-provider-stingray/Godeps/_workspace/src/github.com/hashicorp/terraform/helper/hashcode"
	"github.com/whitepages/terraform-provider-stingray/Godeps/_workspace/src/github.com/hashicorp/terraform/helper/schema"
	"github.com/whitepages/terraform-provider-stingray/Godeps/_workspace/src/github.com/whitepages/go-stingray"
)

func resourceVirtualServer() *schema.Resource {
	return &schema.Resource{
		Create: resourceVirtualServerCreate,
		Read:   resourceVirtualServerRead,
		Update: resourceVirtualServerUpdate,
		Delete: resourceVirtualServerDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"connection_errors_error_file": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"connection_keepalive_timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},

			"connection_timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  10,
			},

			"connect_timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},

			// Default value of enabled in v3.2 of the
			// REST API is false
			"enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},

			"gzip_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"gzip_include_mime": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"http_location_rewrite": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"listen_on_any": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},

			"listen_on_traffic_ips": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"log_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"log_filename": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"log_format": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"log_server_connection_failures": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"note": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"pool": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"port": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
			},

			"protocol": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"recent_connections_save_all": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"request_rules": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},

			"response_rules": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},

			"ssl_add_http_headers": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"ssl_decrypt": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"ssl_server_cert_default": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"ssl_server_cert_host_mapping": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"certificate": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
						"host": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
					},
				},
				Set: func(v interface{}) int {
					m := v.(map[string]interface{})
					return hashcode.String(m["host"].(string))
				},
			},

			"syslog_format": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"web_cache_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"web_cache_max_time": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},
		},
	}
}

func resourceVirtualServerCreate(d *schema.ResourceData, meta interface{}) error {
	err := resourceVirtualServerSet(d, meta)
	if err != nil {
		return err
	}

	return resourceVirtualServerRead(d, meta)
}

func resourceVirtualServerRead(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*providerConfig).client

	r, resp, err := c.GetVirtualServer(d.Get("name").(string))
	if err != nil {
		if resp != nil && resp.StatusCode == 404 {
			// The resource doesn't exist anymore
			d.SetId("")

			return nil
		}

		return fmt.Errorf("Error reading resource: %s", err)
	}

	d.Set("connection_errors_error_file", string(*r.ConnectionErrors.ErrorFile))
	d.Set("connection_keepalive_timeout", int(*r.Connection.KeepaliveTimeout))
	d.Set("connection_timeout", int(*r.Connection.Timeout))
	d.Set("connect_timeout", int(*r.Basic.ConnectTimeout))
	d.Set("enabled", bool(*r.Basic.Enabled))
	d.Set("gzip_enabled", bool(*r.Gzip.Enabled))
	d.Set("gzip_include_mime", []string(*r.Gzip.IncludeMIME))
	d.Set("http_location_rewrite", string(*r.HTTP.LocationRewrite))
	d.Set("listen_on_any", bool(*r.Basic.ListenOnAny))
	d.Set("listen_on_traffic_ips", []string(*r.Basic.ListenOnTrafficIPs))
	d.Set("log_enabled", bool(*r.Log.Enabled))
	d.Set("log_filename", string(*r.Log.Filename))
	d.Set("log_format", string(*r.Log.Format))
	d.Set("log_server_connection_failures", bool(*r.Log.ServerConnectionFailures))
	d.Set("note", string(*r.Basic.Note))
	d.Set("pool", string(*r.Basic.Pool))
	d.Set("port", int(*r.Basic.Port))
	d.Set("protocol", string(*r.Basic.Protocol))
	d.Set("recent_connections_save_all", bool(*r.RecentConnections.SaveAll))
	d.Set("request_rules", []string(*r.Basic.RequestRules))
	d.Set("response_rules", []string(*r.Basic.ResponseRules))
	d.Set("ssl_add_http_headers", bool(*r.SSL.AddHTTPHeaders))
	d.Set("ssl_decrypt", bool(*r.Basic.SSLDecrypt))
	d.Set("ssl_server_cert_default", string(*r.SSL.ServerCertDefault))
	d.Set("ssl_server_cert_host_mapping", flattenServerCertHostMappingTable(*r.SSL.ServerCertHostMapping))
	d.Set("syslog_format", string(*r.Syslog.Format))
	d.Set("web_cache_enabled", bool(*r.WebCache.Enabled))
	d.Set("web_cache_max_time", int(*r.WebCache.MaxTime))

	return nil
}

func resourceVirtualServerUpdate(d *schema.ResourceData, meta interface{}) error {
	err := resourceVirtualServerSet(d, meta)
	if err != nil {
		return err
	}

	return resourceVirtualServerRead(d, meta)
}

func resourceVirtualServerDelete(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*providerConfig).client
	r := stingray.NewVirtualServer(d.Id())

	_, err := c.Delete(r)
	if err != nil {
		return err
	}

	return nil
}

func resourceVirtualServerSet(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*providerConfig).client
	r := stingray.NewVirtualServer(d.Get("name").(string))

	setString(&r.ConnectionErrors.ErrorFile, d, "connection_errors_error_file")
	setInt(&r.Connection.KeepaliveTimeout, d, "connection_keepalive_timeout")
	setInt(&r.Connection.Timeout, d, "connection_timeout")
	setInt(&r.Basic.ConnectTimeout, d, "connect_timeout")
	r.Basic.Enabled = stingray.Bool(d.Get("enabled").(bool))
	setStringSet(&r.Gzip.IncludeMIME, d, "gzip_include_mime")
	setBool(&r.Gzip.Enabled, d, "gzip_enabled")
	setString(&r.HTTP.LocationRewrite, d, "http_location_rewrite")
	r.Basic.ListenOnAny = stingray.Bool(d.Get("listen_on_any").(bool))
	setStringSet(&r.Basic.ListenOnTrafficIPs, d, "listen_on_traffic_ips")
	setBool(&r.Log.Enabled, d, "log_enabled")
	setString(&r.Log.Filename, d, "log_filename")
	setString(&r.Log.Format, d, "log_format")
	setBool(&r.Log.ServerConnectionFailures, d, "log_server_connection_failures")
	setString(&r.Basic.Note, d, "note")
	setString(&r.Basic.Pool, d, "pool")
	setInt(&r.Basic.Port, d, "port")
	setString(&r.Basic.Protocol, d, "protocol")
	setBool(&r.RecentConnections.SaveAll, d, "recent_connections_save_all")
	setStringList(&r.Basic.RequestRules, d, "request_rules")
	setStringList(&r.Basic.ResponseRules, d, "response_rules")
	setBool(&r.SSL.AddHTTPHeaders, d, "ssl_add_http_headers")
	setBool(&r.Basic.SSLDecrypt, d, "ssl_decrypt")
	setString(&r.SSL.ServerCertDefault, d, "ssl_server_cert_default")
	setServerCertHostMappingTable(&r.SSL.ServerCertHostMapping, d, "ssl_server_cert_host_mapping")
	setString(&r.Syslog.Format, d, "syslog_format")
	setBool(&r.WebCache.Enabled, d, "web_cache_enabled")
	setInt(&r.WebCache.MaxTime, d, "web_cache_max_time")

	_, err := c.Set(r)
	if err != nil {
		return err
	}

	d.SetId(d.Get("name").(string))

	return nil
}

func setServerCertHostMappingTable(target **stingray.ServerCertHostMappingTable, d *schema.ResourceData, key string) {
	if _, ok := d.GetOk(key); ok {
		table := d.Get(key).(*schema.Set).List()
		*target, _ = expandServerCertHostMappingTable(table)
	}
}

func expandServerCertHostMappingTable(configured []interface{}) (*stingray.ServerCertHostMappingTable, error) {
	table := make(stingray.ServerCertHostMappingTable, 0, len(configured))

	for _, raw := range configured {
		data := raw.(map[string]interface{})

		s := stingray.ServerCertHostMapping{
			Certificate: stingray.String(data["certificate"].(string)),
			Host:        stingray.String(data["host"].(string)),
		}

		table = append(table, s)
	}

	return &table, nil
}

func flattenServerCertHostMappingTable(list stingray.ServerCertHostMappingTable) []map[string]interface{} {
	result := make([]map[string]interface{}, 0, len(list))
	for _, i := range list {
		s := map[string]interface{}{
			"certificate": *i.Certificate,
			"host":        *i.Host,
		}
		result = append(result, s)
	}

	return result
}
<commit_msg>Add newline<commit_after>package main

import (
	"fmt"

	"github.com/whitepages/terraform-provider-stingray/Godeps/_workspace/src/github.com/hashicorp/terraform/helper/hashcode"
	"github.com/whitepages/terraform-provider-stingray/Godeps/_workspace/src/github.com/hashicorp/terraform/helper/schema"
	"github.com/whitepages/terraform-provider-stingray/Godeps/_workspace/src/github.com/whitepages/go-stingray"
)

func resourceVirtualServer() *schema.Resource {
	return &schema.Resource{
		Create: resourceVirtualServerCreate,
		Read:   resourceVirtualServerRead,
		Update: resourceVirtualServerUpdate,
		Delete: resourceVirtualServerDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"connection_errors_error_file": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"connection_keepalive_timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},

			"connection_timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  10,
			},

			"connect_timeout": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},

			// Default value of enabled in v3.2 of the
			// REST API is false
			"enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},

			"gzip_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"gzip_include_mime": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"http_location_rewrite": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"listen_on_any": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},

			"listen_on_traffic_ips": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Set: func(v interface{}) int {
					return hashcode.String(v.(string))
				},
			},

			"log_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"log_filename": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"log_format": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"log_server_connection_failures": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"note": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"pool": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"port": &schema.Schema{
				Type:     schema.TypeInt,
				Required: true,
			},

			"protocol": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"recent_connections_save_all": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"request_rules": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},

			"response_rules": &schema.Schema{
				Type:     schema.TypeList,
				Optional: true,
				Elem:     &schema.Schema{Type: schema.TypeString},
			},

			"ssl_add_http_headers": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"ssl_decrypt": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"ssl_server_cert_default": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"ssl_server_cert_host_mapping": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Computed: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"certificate": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
						"host": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},
					},
				},
				Set: func(v interface{}) int {
					m := v.(map[string]interface{})
					return hashcode.String(m["host"].(string))
				},
			},

			"syslog_format": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},

			"web_cache_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Computed: true,
			},

			"web_cache_max_time": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},
		},
	}
}

func resourceVirtualServerCreate(d *schema.ResourceData, meta interface{}) error {
	err := resourceVirtualServerSet(d, meta)
	if err != nil {
		return err
	}

	return resourceVirtualServerRead(d, meta)
}

func resourceVirtualServerRead(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*providerConfig).client

	r, resp, err := c.GetVirtualServer(d.Get("name").(string))
	if err != nil {
		if resp != nil && resp.StatusCode == 404 {
			// The resource doesn't exist anymore
			d.SetId("")

			return nil
		}

		return fmt.Errorf("Error reading resource: %s", err)
	}

	d.Set("connection_errors_error_file", string(*r.ConnectionErrors.ErrorFile))
	d.Set("connection_keepalive_timeout", int(*r.Connection.KeepaliveTimeout))
	d.Set("connection_timeout", int(*r.Connection.Timeout))
	d.Set("connect_timeout", int(*r.Basic.ConnectTimeout))
	d.Set("enabled", bool(*r.Basic.Enabled))
	d.Set("gzip_enabled", bool(*r.Gzip.Enabled))
	d.Set("gzip_include_mime", []string(*r.Gzip.IncludeMIME))
	d.Set("http_location_rewrite", string(*r.HTTP.LocationRewrite))
	d.Set("listen_on_any", bool(*r.Basic.ListenOnAny))
	d.Set("listen_on_traffic_ips", []string(*r.Basic.ListenOnTrafficIPs))
	d.Set("log_enabled", bool(*r.Log.Enabled))
	d.Set("log_filename", string(*r.Log.Filename))
	d.Set("log_format", string(*r.Log.Format))
	d.Set("log_server_connection_failures", bool(*r.Log.ServerConnectionFailures))
	d.Set("note", string(*r.Basic.Note))
	d.Set("pool", string(*r.Basic.Pool))
	d.Set("port", int(*r.Basic.Port))
	d.Set("protocol", string(*r.Basic.Protocol))
	d.Set("recent_connections_save_all", bool(*r.RecentConnections.SaveAll))
	d.Set("request_rules", []string(*r.Basic.RequestRules))
	d.Set("response_rules", []string(*r.Basic.ResponseRules))
	d.Set("ssl_add_http_headers", bool(*r.SSL.AddHTTPHeaders))
	d.Set("ssl_decrypt", bool(*r.Basic.SSLDecrypt))
	d.Set("ssl_server_cert_default", string(*r.SSL.ServerCertDefault))
	d.Set("ssl_server_cert_host_mapping", flattenServerCertHostMappingTable(*r.SSL.ServerCertHostMapping))
	d.Set("syslog_format", string(*r.Syslog.Format))
	d.Set("web_cache_enabled", bool(*r.WebCache.Enabled))
	d.Set("web_cache_max_time", int(*r.WebCache.MaxTime))

	return nil
}

func resourceVirtualServerUpdate(d *schema.ResourceData, meta interface{}) error {
	err := resourceVirtualServerSet(d, meta)
	if err != nil {
		return err
	}

	return resourceVirtualServerRead(d, meta)
}

func resourceVirtualServerDelete(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*providerConfig).client
	r := stingray.NewVirtualServer(d.Id())

	_, err := c.Delete(r)
	if err != nil {
		return err
	}

	return nil
}

func resourceVirtualServerSet(d *schema.ResourceData, meta interface{}) error {
	c := meta.(*providerConfig).client
	r := stingray.NewVirtualServer(d.Get("name").(string))

	setString(&r.ConnectionErrors.ErrorFile, d, "connection_errors_error_file")
	setInt(&r.Connection.KeepaliveTimeout, d, "connection_keepalive_timeout")
	setInt(&r.Connection.Timeout, d, "connection_timeout")
	setInt(&r.Basic.ConnectTimeout, d, "connect_timeout")
	r.Basic.Enabled = stingray.Bool(d.Get("enabled").(bool))
	setStringSet(&r.Gzip.IncludeMIME, d, "gzip_include_mime")
	setBool(&r.Gzip.Enabled, d, "gzip_enabled")
	setString(&r.HTTP.LocationRewrite, d, "http_location_rewrite")
	r.Basic.ListenOnAny = stingray.Bool(d.Get("listen_on_any").(bool))
	setStringSet(&r.Basic.ListenOnTrafficIPs, d, "listen_on_traffic_ips")
	setBool(&r.Log.Enabled, d, "log_enabled")
	setString(&r.Log.Filename, d, "log_filename")
	setString(&r.Log.Format, d, "log_format")
	setBool(&r.Log.ServerConnectionFailures, d, "log_server_connection_failures")
	setString(&r.Basic.Note, d, "note")
	setString(&r.Basic.Pool, d, "pool")
	setInt(&r.Basic.Port, d, "port")
	setString(&r.Basic.Protocol, d, "protocol")
	setBool(&r.RecentConnections.SaveAll, d, "recent_connections_save_all")
	setStringList(&r.Basic.RequestRules, d, "request_rules")
	setStringList(&r.Basic.ResponseRules, d, "response_rules")
	setBool(&r.SSL.AddHTTPHeaders, d, "ssl_add_http_headers")
	setBool(&r.Basic.SSLDecrypt, d, "ssl_decrypt")
	setString(&r.SSL.ServerCertDefault, d, "ssl_server_cert_default")
	setServerCertHostMappingTable(&r.SSL.ServerCertHostMapping, d, "ssl_server_cert_host_mapping")
	setString(&r.Syslog.Format, d, "syslog_format")
	setBool(&r.WebCache.Enabled, d, "web_cache_enabled")
	setInt(&r.WebCache.MaxTime, d, "web_cache_max_time")

	_, err := c.Set(r)
	if err != nil {
		return err
	}

	d.SetId(d.Get("name").(string))

	return nil
}

func setServerCertHostMappingTable(target **stingray.ServerCertHostMappingTable, d *schema.ResourceData, key string) {
	if _, ok := d.GetOk(key); ok {
		table := d.Get(key).(*schema.Set).List()
		*target, _ = expandServerCertHostMappingTable(table)
	}
}

func expandServerCertHostMappingTable(configured []interface{}) (*stingray.ServerCertHostMappingTable, error) {
	table := make(stingray.ServerCertHostMappingTable, 0, len(configured))

	for _, raw := range configured {
		data := raw.(map[string]interface{})

		s := stingray.ServerCertHostMapping{
			Certificate: stingray.String(data["certificate"].(string)),
			Host:        stingray.String(data["host"].(string)),
		}

		table = append(table, s)
	}

	return &table, nil
}

func flattenServerCertHostMappingTable(list stingray.ServerCertHostMappingTable) []map[string]interface{} {
	result := make([]map[string]interface{}, 0, len(list))

	for _, i := range list {
		s := map[string]interface{}{
			"certificate": *i.Certificate,
			"host":        *i.Host,
		}
		result = append(result, s)
	}

	return result
}
<|endoftext|>


<commit_before>package v1

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/premkit/premkit/log"
	"github.com/premkit/premkit/models"

	"github.com/vulcand/oxy/forward"
)

var (
	fwd *forward.Forwarder
)

func init() {
	f, err := forward.New()
	if err != nil {
		log.Error(err)
		os.Exit(1)
	}

	fwd = f
}

// TODO Swagger this handler, but it's a special handler and therefore a little trickier to swagger-ify

// ForwardService is the handler for anything that should be possibly fowarded to an upstream.
func ForwardService(response http.ResponseWriter, request *http.Request) {
	// TODO keep these cached because in any reasonable load this will be painful
	services, err := models.ListServices()
	if err != nil {
		http.Error(response, fmt.Sprintf("%+v", err), http.StatusInternalServerError)
		return
	}

	var url *url.URL

	for _, service := range services {
		if !isPathPrefix(service.Path, request.URL.Path) {
			continue
		}

		log.Debugf("path %q matched service %q (service path %q)", request.URL.Path, service.Name, service.Path)

		if len(service.Upstreams) == 0 {
			err := errors.New("No upstreams are available")
			log.Error(err)
			response.WriteHeader(http.StatusBadGateway)
			response.Write([]byte(""))
			return
		}

		childPath := createForwardPath(service.Path, request.URL.Path)
		request.RequestURI = "/" + childPath

		log.Debugf("service.Upstreams = %s/%s", service.Upstreams[0], childPath)
		u, err := url.Parse(fmt.Sprintf("%s/%s", service.Upstreams[0], childPath))
		if err != nil {
			log.Error(err)
			http.Error(response, fmt.Sprintf("%+v", err), http.StatusInternalServerError)
			return
		}

		url = u
	}

	if url == nil {
		http.Error(response, "Route not found", http.StatusNotFound)
		return
	}

	request.URL = url
	fwd.ServeHTTP(response, request)
}

func stripLeadingSlashIfPresent(path string) string {
	return strings.TrimPrefix(path, "/")
}

func isPathPrefix(servicePath, requestPath string) bool {
	servicePath = stripLeadingSlashIfPresent(servicePath)
	requestPath = stripLeadingSlashIfPresent(requestPath)

	return strings.HasPrefix(requestPath, servicePath)
}

func createForwardPath(servicePath, requestPath string) string {
	servicePath = stripLeadingSlashIfPresent(servicePath)
	requestPath = stripLeadingSlashIfPresent(requestPath)

	// Remove the servicePath from the requestPath
	return strings.TrimPrefix(requestPath, servicePath)
}
<commit_msg>No extra slash should be added to the forward<commit_after>package v1

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/premkit/premkit/log"
	"github.com/premkit/premkit/models"

	"github.com/vulcand/oxy/forward"
)

var (
	fwd *forward.Forwarder
)

func init() {
	f, err := forward.New()
	if err != nil {
		log.Error(err)
		os.Exit(1)
	}

	fwd = f
}

// TODO Swagger this handler, but it's a special handler and therefore a little trickier to swagger-ify

// ForwardService is the handler for anything that should be possibly fowarded to an upstream.
func ForwardService(response http.ResponseWriter, request *http.Request) {
	// TODO keep these cached because in any reasonable load this will be painful
	services, err := models.ListServices()
	if err != nil {
		http.Error(response, fmt.Sprintf("%+v", err), http.StatusInternalServerError)
		return
	}

	var url *url.URL

	for _, service := range services {
		if !isPathPrefix(service.Path, request.URL.Path) {
			continue
		}

		log.Debugf("path %q matched service %q (service path %q)", request.URL.Path, service.Name, service.Path)

		if len(service.Upstreams) == 0 {
			err := errors.New("No upstreams are available")
			log.Error(err)
			response.WriteHeader(http.StatusBadGateway)
			response.Write([]byte(""))
			return
		}

		childPath := createForwardPath(service.Path, request.URL.Path)
		request.RequestURI = childPath

		log.Debugf("service.Upstreams = %s/%s", service.Upstreams[0], childPath)
		u, err := url.Parse(fmt.Sprintf("%s/%s", service.Upstreams[0], childPath))
		if err != nil {
			log.Error(err)
			http.Error(response, fmt.Sprintf("%+v", err), http.StatusInternalServerError)
			return
		}

		url = u
	}

	if url == nil {
		http.Error(response, "Route not found", http.StatusNotFound)
		return
	}

	request.URL = url
	fwd.ServeHTTP(response, request)
}

func stripLeadingSlashIfPresent(path string) string {
	return strings.TrimPrefix(path, "/")
}

func isPathPrefix(servicePath, requestPath string) bool {
	servicePath = stripLeadingSlashIfPresent(servicePath)
	requestPath = stripLeadingSlashIfPresent(requestPath)

	return strings.HasPrefix(requestPath, servicePath)
}

func createForwardPath(servicePath, requestPath string) string {
	servicePath = stripLeadingSlashIfPresent(servicePath)
	requestPath = stripLeadingSlashIfPresent(requestPath)

	// Remove the servicePath from the requestPath
	return strings.TrimPrefix(requestPath, servicePath)
}
<|endoftext|>


<commit_before>package bitsgo

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	"github.com/petergtz/bitsgo/logger"
)

type MetricsService interface {
	SendTimingMetric(name string, duration time.Duration)
}

type ResourceHandler struct {
	blobstore        Blobstore
	resourceType     string
	metricsService   MetricsService
	maxBodySizeLimit uint64
}

func NewResourceHandler(blobstore Blobstore, resourceType string, metricsService MetricsService, maxBodySizeLimit uint64) *ResourceHandler {
	return &ResourceHandler{
		blobstore:        blobstore,
		resourceType:     resourceType,
		metricsService:   metricsService,
		maxBodySizeLimit: maxBodySizeLimit,
	}
}

// TODO: instead of params, we could use `identifier string` to make the interface more type-safe.
//       Here and in the other methods.
func (handler *ResourceHandler) PutNoMultiPart(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	if !HandleBodySizeLimits(responseWriter, request, handler.maxBodySizeLimit) {
		return
	}
	logger.From(request).Debugw("Octet-Stream")

	digest := request.Header.Get("Digest")
	if digest == "" {
		badRequest(responseWriter, "No Digest header")
		return
	}
	parts := strings.Split(digest, "=")
	if len(parts) != 2 {
		badRequest(responseWriter, "Digest must have format sha256=value, but is "+digest)
		return
	}
	alg, value := strings.ToLower(parts[0]), parts[1]
	if alg != "sha256" {
		badRequest(responseWriter, "Digest must have format sha256=value, but is "+digest)
		return
	}
	if value == "" {
		badRequest(responseWriter, "Digest must have format sha256=value. Value cannot be empty")
		return
	}

	// TODO this can cause an out of memory panic. Should be smart about writing big files to disk instead.
	content, e := ioutil.ReadAll(request.Body)
	if e != nil {
		internalServerError(responseWriter, e)
		return
	}
	e = handler.blobstore.Put(params["identifier"]+"/"+value, bytes.NewReader(content))
	writeResponseBasedOn("", e, responseWriter, http.StatusCreated, emptyReader)
}

// TODO: instead of params, we could use `identifier string` to make the interface more type-safe.
//       Here and in the other methods.
func (handler *ResourceHandler) AddOrReplace(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	if !HandleBodySizeLimits(responseWriter, request, handler.maxBodySizeLimit) {
		return
	}
	if strings.Contains(request.Header.Get("Content-Type"), "multipart/form-data") {
		logger.From(request).Debugw("Multipart upload")
		handler.uploadMultipart(responseWriter, request, params["identifier"])
	} else {
		logger.From(request).Debugw("Copy source guid")
		handler.copySourceGuid(responseWriter, request, params["identifier"])
	}
}

func (handler *ResourceHandler) uploadMultipart(responseWriter http.ResponseWriter, request *http.Request, identifier string) {
	file, _, e := request.FormFile(handler.resourceType)
	if e != nil {
		badRequest(responseWriter, "Could not retrieve '%s' form parameter", handler.resourceType)
		return
	}
	defer file.Close()

	startTime := time.Now()
	e = handler.blobstore.Put(identifier, file)
	handler.metricsService.SendTimingMetric(handler.resourceType+"-cp_to_blobstore-time", time.Since(startTime))

	writeResponseBasedOn("", e, responseWriter, http.StatusCreated, emptyReader)
}

func (handler *ResourceHandler) copySourceGuid(responseWriter http.ResponseWriter, request *http.Request, identifier string) {
	sourceGuid := sourceGuidFrom(request.Body, responseWriter)
	if sourceGuid == "" {
		return // response is already handled in sourceGuidFrom
	}
	e := handler.blobstore.Copy(sourceGuid, identifier)
	writeResponseBasedOn("", e, responseWriter, http.StatusCreated, emptyReader)
}

func sourceGuidFrom(body io.ReadCloser, responseWriter http.ResponseWriter) string {
	content, e := ioutil.ReadAll(body)
	if e != nil {
		internalServerError(responseWriter, e)
		return ""
	}
	var payload struct {
		SourceGuid string `json:"source_guid"`
	}
	e = json.Unmarshal(content, &payload)
	if e != nil {
		badRequest(responseWriter, "Body must be valid JSON when request is not multipart/form-data. %+v", e)
		return ""
	}
	return payload.SourceGuid
}

func (handler *ResourceHandler) HeadOrRedirectAsGet(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	redirectLocation, e := handler.blobstore.HeadOrRedirectAsGet(params["identifier"])
	writeResponseBasedOn(redirectLocation, e, responseWriter, http.StatusOK, emptyReader)
}

func (handler *ResourceHandler) Get(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	body, redirectLocation, e := handler.blobstore.GetOrRedirect(params["identifier"])
	writeResponseBasedOn(redirectLocation, e, responseWriter, http.StatusOK, body)
}

func (handler *ResourceHandler) Delete(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	// TODO nothing should be S3 specific here
	// this check is needed, because S3 does not return a NotFound on a Delete request:
	exists, e := handler.blobstore.Exists(params["identifier"])
	if e != nil {
		internalServerError(responseWriter, e)
		return
	}
	if !exists {
		responseWriter.WriteHeader(http.StatusNotFound)
		return
	}
	e = handler.blobstore.Delete(params["identifier"])
	writeResponseBasedOn("", e, responseWriter, http.StatusNoContent, emptyReader)
}

func (handler *ResourceHandler) DeleteDir(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	e := handler.blobstore.DeleteDir(params["identifier"])
	switch e.(type) {
	case *NotFoundError:
		responseWriter.WriteHeader(http.StatusNoContent)
		return
	}
	writeResponseBasedOn("", e, responseWriter, http.StatusNoContent, emptyReader)
}

var emptyReader = ioutil.NopCloser(bytes.NewReader(nil))

func writeResponseBasedOn(redirectLocation string, e error, responseWriter http.ResponseWriter, statusCode int, body io.ReadCloser) {
	switch e.(type) {
	case *NotFoundError:
		responseWriter.WriteHeader(http.StatusNotFound)
		return
	case *NoSpaceLeftError:
		responseWriter.WriteHeader(http.StatusInsufficientStorage)
		return
	case error:
		internalServerError(responseWriter, e)
		return
	}
	if redirectLocation != "" {
		redirect(responseWriter, redirectLocation)
		return
	}
	defer body.Close()
	responseWriter.WriteHeader(statusCode)
	io.Copy(responseWriter, body)
}

func redirect(responseWriter http.ResponseWriter, redirectLocation string) {
	responseWriter.Header().Set("Location", redirectLocation)
	responseWriter.WriteHeader(http.StatusFound)
}

func internalServerError(responseWriter http.ResponseWriter, e error) {
	logger.Log.Errorw("Internal Server Error.", "error", fmt.Sprintf("%+v", e))
	responseWriter.WriteHeader(http.StatusInternalServerError)
}

func badRequest(responseWriter http.ResponseWriter, message string, args ...interface{}) {
	responseBody := fmt.Sprintf(message, args...)
	logger.Log.Debugw("Bad request", "body", responseBody)
	responseWriter.WriteHeader(http.StatusBadRequest)
	fprintDescriptionAndCodeAsJSON(responseWriter, "290003", message, args...)
}

func fprintDescriptionAndCodeAsJSON(responseWriter http.ResponseWriter, code string, description string, a ...interface{}) {
	fmt.Fprintf(responseWriter, `{"description":"%v","code":%v}`, fmt.Sprintf(description, a...), code)
}
<commit_msg>Allow form filename 'bits' and return json response<commit_after>package bitsgo

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	"github.com/petergtz/bitsgo/logger"
)

type MetricsService interface {
	SendTimingMetric(name string, duration time.Duration)
}

type ResourceHandler struct {
	blobstore        Blobstore
	resourceType     string
	metricsService   MetricsService
	maxBodySizeLimit uint64
}

func NewResourceHandler(blobstore Blobstore, resourceType string, metricsService MetricsService, maxBodySizeLimit uint64) *ResourceHandler {
	return &ResourceHandler{
		blobstore:        blobstore,
		resourceType:     resourceType,
		metricsService:   metricsService,
		maxBodySizeLimit: maxBodySizeLimit,
	}
}

// TODO: instead of params, we could use `identifier string` to make the interface more type-safe.
//       Here and in the other methods.
func (handler *ResourceHandler) PutNoMultiPart(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	if !HandleBodySizeLimits(responseWriter, request, handler.maxBodySizeLimit) {
		return
	}
	logger.From(request).Debugw("Octet-Stream")

	digest := request.Header.Get("Digest")
	if digest == "" {
		badRequest(responseWriter, "No Digest header")
		return
	}
	parts := strings.Split(digest, "=")
	if len(parts) != 2 {
		badRequest(responseWriter, "Digest must have format sha256=value, but is "+digest)
		return
	}
	alg, value := strings.ToLower(parts[0]), parts[1]
	if alg != "sha256" {
		badRequest(responseWriter, "Digest must have format sha256=value, but is "+digest)
		return
	}
	if value == "" {
		badRequest(responseWriter, "Digest must have format sha256=value. Value cannot be empty")
		return
	}

	// TODO this can cause an out of memory panic. Should be smart about writing big files to disk instead.
	content, e := ioutil.ReadAll(request.Body)
	if e != nil {
		internalServerError(responseWriter, e)
		return
	}
	e = handler.blobstore.Put(params["identifier"]+"/"+value, bytes.NewReader(content))
	writeResponseBasedOn("", e, responseWriter, http.StatusCreated, emptyReader)
}

// TODO: instead of params, we could use `identifier string` to make the interface more type-safe.
//       Here and in the other methods.
func (handler *ResourceHandler) AddOrReplace(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	if !HandleBodySizeLimits(responseWriter, request, handler.maxBodySizeLimit) {
		return
	}
	if strings.Contains(request.Header.Get("Content-Type"), "multipart/form-data") {
		logger.From(request).Debugw("Multipart upload")
		handler.uploadMultipart(responseWriter, request, params["identifier"])
	} else {
		logger.From(request).Debugw("Copy source guid")
		handler.copySourceGuid(responseWriter, request, params["identifier"])
	}
}

func (handler *ResourceHandler) uploadMultipart(responseWriter http.ResponseWriter, request *http.Request, identifier string) {
	file, _, e := request.FormFile(handler.resourceType)
	if e == http.ErrMissingFile {
		file, _, e = request.FormFile("bits")
	}
	if e != nil {
		if e == http.ErrMissingFile {
			badRequest(responseWriter, "Could not retrieve form parameter '%s' or 'bits", handler.resourceType)
		} else {
			internalServerError(responseWriter, e)
		}
		return
	}
	defer file.Close()

	startTime := time.Now()
	e = handler.blobstore.Put(identifier, file)
	handler.metricsService.SendTimingMetric(handler.resourceType+"-cp_to_blobstore-time", time.Since(startTime))

	// TODO use Clock instead:
	respBody, marshallingErr := json.Marshal(responseBody{Guid: identifier, State: "READY", Type: "bits", CreatedAt: time.Now()})

	if marshallingErr != nil {
		internalServerError(responseWriter, marshallingErr)
		return
	}

	writeResponseBasedOn("", e, responseWriter, http.StatusCreated, ioutil.NopCloser(bytes.NewReader(respBody)))
}

type responseBody struct {
	Guid      string    `json:"guid"`
	State     string    `json:"state"`
	Type      string    `json:"type"`
	CreatedAt time.Time `json:"created_at"`
}

func (handler *ResourceHandler) copySourceGuid(responseWriter http.ResponseWriter, request *http.Request, identifier string) {
	sourceGuid := sourceGuidFrom(request.Body, responseWriter)
	if sourceGuid == "" {
		return // response is already handled in sourceGuidFrom
	}
	e := handler.blobstore.Copy(sourceGuid, identifier)
	writeResponseBasedOn("", e, responseWriter, http.StatusCreated, emptyReader)
}

func sourceGuidFrom(body io.ReadCloser, responseWriter http.ResponseWriter) string {
	content, e := ioutil.ReadAll(body)
	if e != nil {
		internalServerError(responseWriter, e)
		return ""
	}
	var payload struct {
		SourceGuid string `json:"source_guid"`
	}
	e = json.Unmarshal(content, &payload)
	if e != nil {
		badRequest(responseWriter, "Body must be valid JSON when request is not multipart/form-data. %+v", e)
		return ""
	}
	return payload.SourceGuid
}

func (handler *ResourceHandler) HeadOrRedirectAsGet(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	redirectLocation, e := handler.blobstore.HeadOrRedirectAsGet(params["identifier"])
	writeResponseBasedOn(redirectLocation, e, responseWriter, http.StatusOK, emptyReader)
}

func (handler *ResourceHandler) Get(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	body, redirectLocation, e := handler.blobstore.GetOrRedirect(params["identifier"])
	writeResponseBasedOn(redirectLocation, e, responseWriter, http.StatusOK, body)
}

func (handler *ResourceHandler) Delete(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	// TODO nothing should be S3 specific here
	// this check is needed, because S3 does not return a NotFound on a Delete request:
	exists, e := handler.blobstore.Exists(params["identifier"])
	if e != nil {
		internalServerError(responseWriter, e)
		return
	}
	if !exists {
		responseWriter.WriteHeader(http.StatusNotFound)
		return
	}
	e = handler.blobstore.Delete(params["identifier"])
	writeResponseBasedOn("", e, responseWriter, http.StatusNoContent, emptyReader)
}

func (handler *ResourceHandler) DeleteDir(responseWriter http.ResponseWriter, request *http.Request, params map[string]string) {
	e := handler.blobstore.DeleteDir(params["identifier"])
	switch e.(type) {
	case *NotFoundError:
		responseWriter.WriteHeader(http.StatusNoContent)
		return
	}
	writeResponseBasedOn("", e, responseWriter, http.StatusNoContent, emptyReader)
}

var emptyReader = ioutil.NopCloser(bytes.NewReader(nil))

func writeResponseBasedOn(redirectLocation string, e error, responseWriter http.ResponseWriter, statusCode int, body io.ReadCloser) {
	switch e.(type) {
	case *NotFoundError:
		responseWriter.WriteHeader(http.StatusNotFound)
		return
	case *NoSpaceLeftError:
		responseWriter.WriteHeader(http.StatusInsufficientStorage)
		return
	case error:
		internalServerError(responseWriter, e)
		return
	}
	if redirectLocation != "" {
		redirect(responseWriter, redirectLocation)
		return
	}
	defer body.Close()
	responseWriter.WriteHeader(statusCode)
	io.Copy(responseWriter, body)
}

func redirect(responseWriter http.ResponseWriter, redirectLocation string) {
	responseWriter.Header().Set("Location", redirectLocation)
	responseWriter.WriteHeader(http.StatusFound)
}

func internalServerError(responseWriter http.ResponseWriter, e error) {
	logger.Log.Errorw("Internal Server Error.", "error", fmt.Sprintf("%+v", e))
	responseWriter.WriteHeader(http.StatusInternalServerError)
}

func badRequest(responseWriter http.ResponseWriter, message string, args ...interface{}) {
	responseBody := fmt.Sprintf(message, args...)
	logger.Log.Infow("Bad request", "body", responseBody)
	responseWriter.WriteHeader(http.StatusBadRequest)
	fprintDescriptionAndCodeAsJSON(responseWriter, "290003", message, args...)
}

func fprintDescriptionAndCodeAsJSON(responseWriter http.ResponseWriter, code string, description string, a ...interface{}) {
	fmt.Fprintf(responseWriter, `{"description":"%v","code":%v}`, fmt.Sprintf(description, a...), code)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
)

type Member struct {
	Name     string
	Addr     string
	Port     int
	Tags     map[string]string
	Status   string
	Protocol map[string]int
}

type SerfOutput struct {
	Members []Member
}

func ParseMembers(serf_output []byte) ([]Member, error) {
	var result SerfOutput
	err := json.Unmarshal(serf_output, &result)
	if err != nil {
		return nil, err
	}
	return result.Members, nil
}
<commit_msg>strip port in address<commit_after>package main

import (
	"encoding/json"
	"strconv"
	"strings"
)

type Member struct {
	Name     string
	Addr     string
	Port     int
	Tags     map[string]string
	Status   string
	Protocol map[string]int
}

type SerfOutput struct {
	Members []Member
}

func ParseMembers(serf_output []byte) ([]Member, error) {
	var result SerfOutput
	err := json.Unmarshal(serf_output, &result)
	if err != nil {
		return nil, err
	}
	for i := 0; i < len(result.Members); i++ {
		result.Members[i].Addr = strings.Split(result.Members[i].Addr, ":"+strconv.Itoa(result.Members[i].Port))[0]
	}
	return result.Members, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gin-gonic/contrib/sessions"
	"github.com/gin-gonic/gin"
	"github.com/pariz/gountries"
	"github.com/rjeczalik/notify"
	"github.com/thehowl/conf"
	"zxq.co/ripple/rippleapi/common"
)

var templates = make(map[string]*template.Template)
var baseTemplates = [...]string{
	"templates/base.html",
	"templates/navbar.html",
	"templates/simplepag.html",
}
var simplePages []templateConfig

var gdb = gountries.New()

func countryReadable(s string) string {
	if s == "XX" || s == "" {
		return ""
	}
	reg, err := gdb.FindCountryByAlpha(s)
	if err != nil {
		return ""
	}
	return reg.Name.Common
}

func loadTemplates(subdir string) {
	ts, err := ioutil.ReadDir("templates" + subdir)
	if err != nil {
		panic(err)
	}

	for _, i := range ts {
		// if it's a directory, load recursively
		if i.IsDir() && i.Name() != ".." && i.Name() != "." {
			loadTemplates(subdir + "/" + i.Name())
			continue
		}

		fullName := "templates" + subdir + "/" + i.Name()
		_c := parseConfig(fullName)
		var c templateConfig
		if _c != nil {
			c = *_c
		}
		if c.NoCompile {
			continue
		}

		var files = c.inc("templates" + subdir + "/")
		files = append(files, fullName)

		// do not compile base templates on their own
		var comp bool
		for _, j := range baseTemplates {
			if fullName == j {
				comp = true
				break
			}
		}
		if comp {
			continue
		}

		var inName string
		if subdir != "" && subdir[0] == '/' {
			inName = subdir[1:] + "/"
		}

		// add new template to template slice
		templates[inName+i.Name()] = template.Must(template.New(i.Name()).Funcs(funcMap).ParseFiles(
			append(files, baseTemplates[:]...)...,
		))

		if _c != nil {
			simplePages = append(simplePages, *_c)
		}
	}
}

func resp(c *gin.Context, statusCode int, tpl string, data interface{}) {
	if c == nil {
		return
	}
	t := templates[tpl]
	if t == nil {
		c.String(500, "Template not found! Please tell this to a dev!")
		return
	}
	sess := getSession(c)
	if corrected, ok := data.(page); ok {
		corrected.SetMessages(getMessages(c))
		corrected.SetPath(c.Request.URL.Path)
		corrected.SetContext(getContext(c))
		corrected.SetGinContext(c)
		corrected.SetSession(sess)
	}
	sess.Save()
	buf := &bytes.Buffer{}
	err := t.ExecuteTemplate(buf, "base", data)
	if err != nil {
		c.String(
			200,
			"An error occurred while trying to render the page, and we have now been notified about it.",
		)
		c.Error(err)
		return
	}
	c.Header("Content-Type", "text/html; charset=utf-8")
	c.Status(statusCode)
	_, err = io.Copy(c.Writer, buf)
	if err != nil {
		c.Writer.WriteString("We don't know what's happening now.")
		c.Error(err)
		return
	}
}

type baseTemplateData struct {
	TitleBar          string // required
	HeadingTitle      string
	HeadingOnRight    bool
	Scripts           []string
	KyutGrill         string
	KyutGrillAbsolute bool
	SolidColour       string
	DisableHH         bool // HH = Huge Heading
	Messages          []message
	RequestInfo       map[string]interface{}

	// ignore, they're set by resp()
	Context  context
	Path     string
	FormData map[string]string
	Gin      *gin.Context
	Session  sessions.Session
}

func (b *baseTemplateData) SetMessages(m []message) {
	b.Messages = append(b.Messages, m...)
}
func (b *baseTemplateData) SetPath(path string) {
	b.Path = path
}
func (b *baseTemplateData) SetContext(c context) {
	b.Context = c
}
func (b *baseTemplateData) SetGinContext(c *gin.Context) {
	b.Gin = c
}
func (b *baseTemplateData) SetSession(sess sessions.Session) {
	b.Session = sess
}
func (b baseTemplateData) Get(s string, params ...interface{}) map[string]interface{} {
	s = fmt.Sprintf(s, params...)
	req, err := http.NewRequest("GET", config.API+s, nil)
	if err != nil {
		b.Gin.Error(err)
		return nil
	}
	req.Header.Set("User-Agent", "hanayo")
	req.Header.Set("H-Key", config.APISecret)
	req.Header.Set("X-Ripple-Token", b.Context.Token)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		b.Gin.Error(err)
		return nil
	}
	data, err := ioutil.ReadAll(resp.Body)
	defer resp.Body.Close()
	if err != nil {
		b.Gin.Error(err)
		return nil
	}
	x := make(map[string]interface{})
	err = json.Unmarshal(data, &x)
	if err != nil {
		b.Gin.Error(err)
		return nil
	}
	return x
}
func (b baseTemplateData) Has(privs uint64) bool {
	return uint64(b.Context.User.Privileges)&privs == privs
}
func (b baseTemplateData) Conf() interface{} {
	return config
}

// list of client flags
const (
	CFDarkSite = 1 << iota
)

func (b baseTemplateData) ClientFlags() int {
	s, _ := b.Gin.Cookie("cflags")
	return common.Int(s)
}

type page interface {
	SetMessages([]message)
	SetPath(string)
	SetContext(context)
	SetGinContext(*gin.Context)
	SetSession(sessions.Session)
}

func reloader() error {
	c := make(chan notify.EventInfo, 1)
	if err := notify.Watch("./templates/...", c, notify.All); err != nil {
		return err
	}
	go func() {
		var last time.Time
		for range c {
			if time.Since(last) < time.Second*3 {
				continue
			}
			fmt.Println("Change detected! Refreshing templates")
			simplePages = []templateConfig{}
			loadTemplates("")
			l.Close()
			last = time.Now()
		}
		defer notify.Stop(c)
	}()
	return nil
}

type templateConfig struct {
	NoCompile bool
	Include   string
	Template  string

	// Stuff that used to be in simpleTemplate
	Handler          string
	TitleBar         string
	KyutGrill        string
	MinPrivileges    uint64
	HugeHeadingRight bool
	AdditionalJS     string
}

func (t templateConfig) inc(prefix string) []string {
	if t.Include == "" {
		return nil
	}
	a := strings.Split(t.Include, ",")
	for i, s := range a {
		a[i] = prefix + s
	}
	return a
}

func (t templateConfig) mp() common.UserPrivileges {
	return common.UserPrivileges(t.MinPrivileges)
}

func (t templateConfig) additionalJS() []string {
	parts := strings.Split(t.AdditionalJS, ",")
	if len(parts) > 0 && parts[len(parts)-1] == "" {
		parts = parts[:len(parts)-1]
	}
	return parts
}

func parseConfig(s string) *templateConfig {
	f, err := os.Open(s)
	defer f.Close()
	if err != nil {
		return nil
	}
	i := bufio.NewScanner(f)
	var inConfig bool
	var buff string
	var t templateConfig
	for i.Scan() {
		u := i.Text()
		switch u {
		case "{{/*###":
			inConfig = true
		case "*/}}":
			if !inConfig {
				continue
			}
			conf.LoadRaw(&t, []byte(buff))
			t.Template = strings.TrimPrefix(s, "templates/")
			return &t
		}
		if !inConfig {
			continue
		}
		buff += u + "\n"
	}
	return nil
}

func respEmpty(c *gin.Context, title string, messages ...message) {
	resp(c, 200, "empty.html", &baseTemplateData{TitleBar: title, Messages: messages})
}
<commit_msg>Only reload template changes for HTML files<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/gin-gonic/contrib/sessions"
	"github.com/gin-gonic/gin"
	"github.com/pariz/gountries"
	"github.com/rjeczalik/notify"
	"github.com/thehowl/conf"
	"zxq.co/ripple/rippleapi/common"
)

var templates = make(map[string]*template.Template)
var baseTemplates = [...]string{
	"templates/base.html",
	"templates/navbar.html",
	"templates/simplepag.html",
}
var simplePages []templateConfig

var gdb = gountries.New()

func countryReadable(s string) string {
	if s == "XX" || s == "" {
		return ""
	}
	reg, err := gdb.FindCountryByAlpha(s)
	if err != nil {
		return ""
	}
	return reg.Name.Common
}

func loadTemplates(subdir string) {
	ts, err := ioutil.ReadDir("templates" + subdir)
	if err != nil {
		panic(err)
	}

	for _, i := range ts {
		// if it's a directory, load recursively
		if i.IsDir() && i.Name() != ".." && i.Name() != "." {
			loadTemplates(subdir + "/" + i.Name())
			continue
		}

		fullName := "templates" + subdir + "/" + i.Name()
		_c := parseConfig(fullName)
		var c templateConfig
		if _c != nil {
			c = *_c
		}
		if c.NoCompile {
			continue
		}

		var files = c.inc("templates" + subdir + "/")
		files = append(files, fullName)

		// do not compile base templates on their own
		var comp bool
		for _, j := range baseTemplates {
			if fullName == j {
				comp = true
				break
			}
		}
		if comp {
			continue
		}

		var inName string
		if subdir != "" && subdir[0] == '/' {
			inName = subdir[1:] + "/"
		}

		// add new template to template slice
		templates[inName+i.Name()] = template.Must(template.New(i.Name()).Funcs(funcMap).ParseFiles(
			append(files, baseTemplates[:]...)...,
		))

		if _c != nil {
			simplePages = append(simplePages, *_c)
		}
	}
}

func resp(c *gin.Context, statusCode int, tpl string, data interface{}) {
	if c == nil {
		return
	}
	t := templates[tpl]
	if t == nil {
		c.String(500, "Template not found! Please tell this to a dev!")
		return
	}
	sess := getSession(c)
	if corrected, ok := data.(page); ok {
		corrected.SetMessages(getMessages(c))
		corrected.SetPath(c.Request.URL.Path)
		corrected.SetContext(getContext(c))
		corrected.SetGinContext(c)
		corrected.SetSession(sess)
	}
	sess.Save()
	buf := &bytes.Buffer{}
	err := t.ExecuteTemplate(buf, "base", data)
	if err != nil {
		c.String(
			200,
			"An error occurred while trying to render the page, and we have now been notified about it.",
		)
		c.Error(err)
		return
	}
	c.Header("Content-Type", "text/html; charset=utf-8")
	c.Status(statusCode)
	_, err = io.Copy(c.Writer, buf)
	if err != nil {
		c.Writer.WriteString("We don't know what's happening now.")
		c.Error(err)
		return
	}
}

type baseTemplateData struct {
	TitleBar          string // required
	HeadingTitle      string
	HeadingOnRight    bool
	Scripts           []string
	KyutGrill         string
	KyutGrillAbsolute bool
	SolidColour       string
	DisableHH         bool // HH = Huge Heading
	Messages          []message
	RequestInfo       map[string]interface{}

	// ignore, they're set by resp()
	Context  context
	Path     string
	FormData map[string]string
	Gin      *gin.Context
	Session  sessions.Session
}

func (b *baseTemplateData) SetMessages(m []message) {
	b.Messages = append(b.Messages, m...)
}
func (b *baseTemplateData) SetPath(path string) {
	b.Path = path
}
func (b *baseTemplateData) SetContext(c context) {
	b.Context = c
}
func (b *baseTemplateData) SetGinContext(c *gin.Context) {
	b.Gin = c
}
func (b *baseTemplateData) SetSession(sess sessions.Session) {
	b.Session = sess
}
func (b baseTemplateData) Get(s string, params ...interface{}) map[string]interface{} {
	s = fmt.Sprintf(s, params...)
	req, err := http.NewRequest("GET", config.API+s, nil)
	if err != nil {
		b.Gin.Error(err)
		return nil
	}
	req.Header.Set("User-Agent", "hanayo")
	req.Header.Set("H-Key", config.APISecret)
	req.Header.Set("X-Ripple-Token", b.Context.Token)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		b.Gin.Error(err)
		return nil
	}
	data, err := ioutil.ReadAll(resp.Body)
	defer resp.Body.Close()
	if err != nil {
		b.Gin.Error(err)
		return nil
	}
	x := make(map[string]interface{})
	err = json.Unmarshal(data, &x)
	if err != nil {
		b.Gin.Error(err)
		return nil
	}
	return x
}
func (b baseTemplateData) Has(privs uint64) bool {
	return uint64(b.Context.User.Privileges)&privs == privs
}
func (b baseTemplateData) Conf() interface{} {
	return config
}

// list of client flags
const (
	CFDarkSite = 1 << iota
)

func (b baseTemplateData) ClientFlags() int {
	s, _ := b.Gin.Cookie("cflags")
	return common.Int(s)
}

type page interface {
	SetMessages([]message)
	SetPath(string)
	SetContext(context)
	SetGinContext(*gin.Context)
	SetSession(sessions.Session)
}

func reloader() error {
	c := make(chan notify.EventInfo, 1)
	if err := notify.Watch("./templates/...", c, notify.All); err != nil {
		return err
	}
	go func() {
		var last time.Time
		for ev := range c {
			if !strings.HasSuffix(ev.Path(), ".html") || time.Since(last) < time.Second*3 {
				continue
			}
			fmt.Println("Change detected! Refreshing templates")
			simplePages = []templateConfig{}
			loadTemplates("")
			l.Close()
			last = time.Now()
		}
		defer notify.Stop(c)
	}()
	return nil
}

type templateConfig struct {
	NoCompile bool
	Include   string
	Template  string

	// Stuff that used to be in simpleTemplate
	Handler          string
	TitleBar         string
	KyutGrill        string
	MinPrivileges    uint64
	HugeHeadingRight bool
	AdditionalJS     string
}

func (t templateConfig) inc(prefix string) []string {
	if t.Include == "" {
		return nil
	}
	a := strings.Split(t.Include, ",")
	for i, s := range a {
		a[i] = prefix + s
	}
	return a
}

func (t templateConfig) mp() common.UserPrivileges {
	return common.UserPrivileges(t.MinPrivileges)
}

func (t templateConfig) additionalJS() []string {
	parts := strings.Split(t.AdditionalJS, ",")
	if len(parts) > 0 && parts[len(parts)-1] == "" {
		parts = parts[:len(parts)-1]
	}
	return parts
}

func parseConfig(s string) *templateConfig {
	f, err := os.Open(s)
	defer f.Close()
	if err != nil {
		return nil
	}
	i := bufio.NewScanner(f)
	var inConfig bool
	var buff string
	var t templateConfig
	for i.Scan() {
		u := i.Text()
		switch u {
		case "{{/*###":
			inConfig = true
		case "*/}}":
			if !inConfig {
				continue
			}
			conf.LoadRaw(&t, []byte(buff))
			t.Template = strings.TrimPrefix(s, "templates/")
			return &t
		}
		if !inConfig {
			continue
		}
		buff += u + "\n"
	}
	return nil
}

func respEmpty(c *gin.Context, title string, messages ...message) {
	resp(c, 200, "empty.html", &baseTemplateData{TitleBar: title, Messages: messages})
}
<|endoftext|>


<commit_before>package xlsx

import (
	"fmt"
	"regexp"
	"text/template"
	"time"
)

// Templates for the various XML content in XLST files
var (
	TemplateContentTypes          *template.Template
	TemplateRelationships         *template.Template
	TemplateWorkbook              *template.Template
	TemplateWorkbookRelationships *template.Template
	TemplateStyles                *template.Template
	TemplateStringLookups         *template.Template
	TemplateSheetStart            *template.Template
	TemplateSheetEnd              *template.Template
	TemplateCellNumber            *template.Template
	TemplateCellString            *template.Template
	TemplateCellDateTime          *template.Template
	TemplateCellInlineString      *template.Template
	TemplateApp                   *template.Template
	TemplateCore                  *template.Template
)

// Template function for integer addition. This is useful to convert between
// zero-based and one-based array offsets within templates
func plus(i int, n int) string {
	return fmt.Sprintf("%d", i+n)
}

// Template function for time formatting
func timeFormat(t time.Time) string {
	return t.Format(time.RFC3339)
}

func init() {
	re := regexp.MustCompile("\n[\t\n\f\r ]*")
	funcMap := template.FuncMap{"plus": plus, "timeFormat": timeFormat}

	TemplateContentTypes = template.Must(template.New("templateContentTypes").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateContentTypes, "")))
	TemplateRelationships = template.Must(template.New("templateRelationships").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateRelationships, "")))
	TemplateWorkbook = template.Must(template.New("templateWorkbook").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateWorkbook, "")))
	TemplateWorkbookRelationships = template.Must(template.New("templateWorkbookRelationships").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateWorkbookRelationships, "")))
	TemplateStyles = template.Must(template.New("templateStyles").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateStyles, "")))
	TemplateStringLookups = template.Must(template.New("templateStringLookups").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateStringLookups, "")))
	TemplateSheetStart = template.Must(template.New("templateSheetStart").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateSheetStart, "")))
	TemplateSheetEnd = template.Must(template.New("templateSheetEnd").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateSheetEnd, "")))
	TemplateCellNumber = template.Must(template.New("templateCellNumber").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCellNumber, "")))
	TemplateCellString = template.Must(template.New("templateCellString").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCellString, "")))
	TemplateCellDateTime = template.Must(template.New("templateCellDateTime").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCellDateTime, "")))
	TemplateCellInlineString = template.Must(template.New("templateCellInlineString").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCellInlineString, "")))
	TemplateApp = template.Must(template.New("templateApp").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateApp, "")))
	TemplateCore = template.Must(template.New("templateCore").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCore, "")))
}

const templateContentTypes = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
      <Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>
      <Default Extension="xml" ContentType="application/xml"/>
      <Override PartName="/xl/workbook.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>
      <Override PartName="/xl/worksheets/sheet1.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>
      <Override PartName="/xl/styles.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.styles+xml"/>
      <Override PartName="/xl/sharedStrings.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sharedStrings+xml"/>
      <Override PartName="/docProps/core.xml" ContentType="application/vnd.openxmlformats-package.core-properties+xml"/>
      <Override PartName="/docProps/app.xml" ContentType="application/vnd.openxmlformats-officedocument.extended-properties+xml"/>
  </Types>`

const templateRelationships = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
      <Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="xl/workbook.xml"/>
      <Relationship Id="rId3" Type="http://schemas.openxmlformats.org/package/2006/relationships/metadata/core-properties" Target="docProps/core.xml"/>
      <Relationship Id="rId4" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/extended-properties" Target="docProps/app.xml"/>
  </Relationships>`

const templateWorkbook = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">
      <fileVersion appName="xl" lastEdited="5" lowestEdited="5" rupBuild="9303"/>
      <workbookPr defaultThemeVersion="124226"/>
      <bookViews>
          <workbookView xWindow="480" yWindow="60" windowWidth="18195" windowHeight="8505"/>
      </bookViews>
      <sheets>
          <sheet name="{{.Title}}" sheetId="1" r:id="rId1"/>
      </sheets>
      <calcPr calcId="145621"/>
  </workbook>`

const templateWorkbookRelationships = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
      <Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" Target="worksheets/sheet1.xml"/>
      <Relationship Id="rId2" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/sharedStrings" Target="sharedStrings.xml"/>
      <Relationship Id="rId3" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/styles" Target="styles.xml"/>
  </Relationships>`

const templateStyles = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  `

const templateStringLookups = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<sst xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" count="{{len .}}" uniqueCount="{{len .}}">
{{range .}}<si><t>{{.}}</t></si>{{end}}
</sst>`

const templateSheetStart = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships" xmlns:mc="http://schemas.openxmlformats.org/markup-compatibility/2006" mc:Ignorable="x14ac" xmlns:x14ac="http://schemas.microsoft.com/office/spreadsheetml/2009/9/ac">
            <sheetViews>
        <sheetView workbookViewId="0"/>
      </sheetViews>
      <sheetFormatPr defaultRowHeight="15" x14ac:dyDescent="0.25"/>
        <cols>
          {{range $i, $e := .Cols}}
          <col min="{{plus $i 1}}" max="{{plus $i 1}}" width="{{$e.Width}}" customWidth="1" style="1"/>
          {{end}}
        </cols>
      <sheetData>`

const templateSheetEnd = `
 </sheetData>
 <dimension ref="{{.Start}}:{{.End}}"/>
   </worksheet>`

const templateCellNumber = `<c r="{{.CellIndex}}" t="n" s="1"><v>{{.Value}}</v></c>`
const templateCellString = `<c r="{{.CellIndex}}" t="s" s="1"><v>{{.Value}}</v></c>`
const templateCellDateTime = `<c r="{{.CellIndex}}" s="2"><v>{{.Value}}</v></c>`
const templateCellInlineString = `<c r="{{.CellIndex}}" t="inlineStr"><is><t>{{.Value}}</t></is></c>`

const templateApp = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <Properties xmlns="http://schemas.openxmlformats.org/officeDocument/2006/extended-properties" xmlns:vt="http://schemas.openxmlformats.org/officeDocument/2006/docPropsVTypes">
  <Application>None</Application>
  <DocSecurity>0</DocSecurity>
  <ScaleCrop>false</ScaleCrop>
  <HeadingPairs>
    <vt:vector size="2" baseType="variant">
      <vt:variant>
        <vt:lpstr>Worksheets</vt:lpstr>
      </vt:variant>
      <vt:variant>
        <vt:i4>1</vt:i4>
      </vt:variant>
    </vt:vector>
  </HeadingPairs>
  <TitlesOfParts>
    <vt:vector size="1" baseType="lpstr">
      <vt:lpstr>{{.Title}}</vt:lpstr>
    </vt:vector>
  </TitlesOfParts>
  <LinksUpToDate>false</LinksUpToDate>
  <SharedDoc>false</SharedDoc>
  <HyperlinksChanged>false</HyperlinksChanged>
</Properties>`

const templateCore = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <cp:coreProperties xmlns:cp="http://schemas.openxmlformats.org/package/2006/metadata/core-properties" xmlns:dc="http://purl.org/dc/elements/1.1/" xmlns:dcterms="http://purl.org/dc/terms/" xmlns:dcmitype="http://purl.org/dc/dcmitype/" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
    <dc:creator>{{.CreatedBy}}</dc:creator>
    <cp:lastModifiedBy>{{.ModifiedBy}}</cp:lastModifiedBy>
    <dcterms:created xsi:type="dcterms:W3CDTF">{{timeFormat .CreatedAt}}</dcterms:created>
    <dcterms:modified xsi:type="dcterms:W3CDTF">{{timeFormat .ModifiedAt}}</dcterms:modified>
  </cp:coreProperties>`
<commit_msg>Moved dimension back to sheet level<commit_after>package xlsx

import (
	"fmt"
	"regexp"
	"text/template"
	"time"
)

// Templates for the various XML content in XLST files
var (
	TemplateContentTypes          *template.Template
	TemplateRelationships         *template.Template
	TemplateWorkbook              *template.Template
	TemplateWorkbookRelationships *template.Template
	TemplateStyles                *template.Template
	TemplateStringLookups         *template.Template
	TemplateSheetStart            *template.Template
	TemplateSheetEnd              *template.Template
	TemplateCellNumber            *template.Template
	TemplateCellString            *template.Template
	TemplateCellDateTime          *template.Template
	TemplateCellInlineString      *template.Template
	TemplateApp                   *template.Template
	TemplateCore                  *template.Template
)

// Template function for integer addition. This is useful to convert between
// zero-based and one-based array offsets within templates
func plus(i int, n int) string {
	return fmt.Sprintf("%d", i+n)
}

// Template function for time formatting
func timeFormat(t time.Time) string {
	return t.Format(time.RFC3339)
}

func init() {
	re := regexp.MustCompile("\n[\t\n\f\r ]*")
	funcMap := template.FuncMap{"plus": plus, "timeFormat": timeFormat}

	TemplateContentTypes = template.Must(template.New("templateContentTypes").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateContentTypes, "")))
	TemplateRelationships = template.Must(template.New("templateRelationships").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateRelationships, "")))
	TemplateWorkbook = template.Must(template.New("templateWorkbook").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateWorkbook, "")))
	TemplateWorkbookRelationships = template.Must(template.New("templateWorkbookRelationships").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateWorkbookRelationships, "")))
	TemplateStyles = template.Must(template.New("templateStyles").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateStyles, "")))
	TemplateStringLookups = template.Must(template.New("templateStringLookups").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateStringLookups, "")))
	TemplateSheetStart = template.Must(template.New("templateSheetStart").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateSheetStart, "")))
	TemplateSheetEnd = template.Must(template.New("templateSheetEnd").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateSheetEnd, "")))
	TemplateCellNumber = template.Must(template.New("templateCellNumber").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCellNumber, "")))
	TemplateCellString = template.Must(template.New("templateCellString").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCellString, "")))
	TemplateCellDateTime = template.Must(template.New("templateCellDateTime").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCellDateTime, "")))
	TemplateCellInlineString = template.Must(template.New("templateCellInlineString").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCellInlineString, "")))
	TemplateApp = template.Must(template.New("templateApp").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateApp, "")))
	TemplateCore = template.Must(template.New("templateCore").Funcs(funcMap).Parse(re.ReplaceAllLiteralString(templateCore, "")))
}

const templateContentTypes = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
      <Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>
      <Default Extension="xml" ContentType="application/xml"/>
      <Override PartName="/xl/workbook.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>
      <Override PartName="/xl/worksheets/sheet1.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>
      <Override PartName="/xl/styles.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.styles+xml"/>
      <Override PartName="/xl/sharedStrings.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sharedStrings+xml"/>
      <Override PartName="/docProps/core.xml" ContentType="application/vnd.openxmlformats-package.core-properties+xml"/>
      <Override PartName="/docProps/app.xml" ContentType="application/vnd.openxmlformats-officedocument.extended-properties+xml"/>
  </Types>`

const templateRelationships = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
      <Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="xl/workbook.xml"/>
      <Relationship Id="rId3" Type="http://schemas.openxmlformats.org/package/2006/relationships/metadata/core-properties" Target="docProps/core.xml"/>
      <Relationship Id="rId4" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/extended-properties" Target="docProps/app.xml"/>
  </Relationships>`

const templateWorkbook = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">
      <fileVersion appName="xl" lastEdited="5" lowestEdited="5" rupBuild="9303"/>
      <workbookPr defaultThemeVersion="124226"/>
      <bookViews>
          <workbookView xWindow="480" yWindow="60" windowWidth="18195" windowHeight="8505"/>
      </bookViews>
      <sheets>
          <sheet name="{{.Title}}" sheetId="1" r:id="rId1"/>
      </sheets>
      <calcPr calcId="145621"/>
  </workbook>`

const templateWorkbookRelationships = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
      <Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" Target="worksheets/sheet1.xml"/>
      <Relationship Id="rId2" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/sharedStrings" Target="sharedStrings.xml"/>
      <Relationship Id="rId3" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/styles" Target="styles.xml"/>
  </Relationships>`

const templateStyles = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  `

const templateStringLookups = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<sst xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" count="{{len .}}" uniqueCount="{{len .}}">
{{range .}}<si><t>{{.}}</t></si>{{end}}
</sst>`

const templateSheetStart = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships" xmlns:mc="http://schemas.openxmlformats.org/markup-compatibility/2006" mc:Ignorable="x14ac" xmlns:x14ac="http://schemas.microsoft.com/office/spreadsheetml/2009/9/ac">
            <sheetViews>
        <sheetView workbookViewId="0"/>
      </sheetViews>
      <sheetFormatPr defaultRowHeight="15" x14ac:dyDescent="0.25"/>
        <cols>
          {{range $i, $e := .Cols}}
          <col min="{{plus $i 1}}" max="{{plus $i 1}}" width="{{$e.Width}}" customWidth="1" style="1"/>
          {{end}}
        </cols>
      <sheetData>`

const templateSheetEnd = `
    <dimension ref="{{.Start}}:{{.End}}"/>
    </sheetData>
  </worksheet>`

const templateCellNumber = `<c r="{{.CellIndex}}" t="n" s="1"><v>{{.Value}}</v></c>`
const templateCellString = `<c r="{{.CellIndex}}" t="s" s="1"><v>{{.Value}}</v></c>`
const templateCellDateTime = `<c r="{{.CellIndex}}" s="2"><v>{{.Value}}</v></c>`
const templateCellInlineString = `<c r="{{.CellIndex}}" t="inlineStr"><is><t>{{.Value}}</t></is></c>`

const templateApp = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <Properties xmlns="http://schemas.openxmlformats.org/officeDocument/2006/extended-properties" xmlns:vt="http://schemas.openxmlformats.org/officeDocument/2006/docPropsVTypes">
  <Application>None</Application>
  <DocSecurity>0</DocSecurity>
  <ScaleCrop>false</ScaleCrop>
  <HeadingPairs>
    <vt:vector size="2" baseType="variant">
      <vt:variant>
        <vt:lpstr>Worksheets</vt:lpstr>
      </vt:variant>
      <vt:variant>
        <vt:i4>1</vt:i4>
      </vt:variant>
    </vt:vector>
  </HeadingPairs>
  <TitlesOfParts>
    <vt:vector size="1" baseType="lpstr">
      <vt:lpstr>{{.Title}}</vt:lpstr>
    </vt:vector>
  </TitlesOfParts>
  <LinksUpToDate>false</LinksUpToDate>
  <SharedDoc>false</SharedDoc>
  <HyperlinksChanged>false</HyperlinksChanged>
</Properties>`

const templateCore = `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
  <cp:coreProperties xmlns:cp="http://schemas.openxmlformats.org/package/2006/metadata/core-properties" xmlns:dc="http://purl.org/dc/elements/1.1/" xmlns:dcterms="http://purl.org/dc/terms/" xmlns:dcmitype="http://purl.org/dc/dcmitype/" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
    <dc:creator>{{.CreatedBy}}</dc:creator>
    <cp:lastModifiedBy>{{.ModifiedBy}}</cp:lastModifiedBy>
    <dcterms:created xsi:type="dcterms:W3CDTF">{{timeFormat .CreatedAt}}</dcterms:created>
    <dcterms:modified xsi:type="dcterms:W3CDTF">{{timeFormat .ModifiedAt}}</dcterms:modified>
  </cp:coreProperties>`
<|endoftext|>


<commit_before>package easyss

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/miekg/dns"
	"github.com/nange/easypool"
	"github.com/nange/easyss/cipherstream"
	"github.com/nange/easyss/util"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"github.com/txthinking/socks5"
)

const (
	MaxUDPDataSize   = 65507
	DefaultDNSServer = "8.8.8.8:53"
)

var udpDataBytes = util.NewBytes(MaxUDPDataSize)

// UDPExchange used to store client address and remote connection
type UDPExchange struct {
	ClientAddr *net.UDPAddr
	RemoteConn net.Conn
}

func (ss *Easyss) UDPHandle(s *socks5.Server, addr *net.UDPAddr, d *socks5.Datagram) error {
	log.Debugf("enter udp handdle, local_addr:%v, remote_addr:%v", addr.String(), d.Address())

	dst := d.Address()
	rewrittenDst := dst

	msg := &dns.Msg{}
	err := msg.Unpack(d.Data)
	if err == nil && isDNSRequest(msg) {
		log.Infof("the udp request is dns proto, domain:%s, qtype:%s",
			msg.Question[0].Name, dns.TypeToString[msg.Question[0].Qtype])

		question := msg.Question[0]
		hostAtCN := ss.HostAtCN(strings.TrimSuffix(question.Name, "."))
		isDirect := hostAtCN && ss.Tun2socksStatusAuto()

		// find from dns cache first
		msgCache := ss.DNSCache(question.Name, dns.TypeToString[question.Qtype], isDirect)
		if msgCache != nil {
			msgCache.MsgHdr.Id = msg.MsgHdr.Id
			log.Infof("find msg from dns cache, write back directly, domain:%s, qtype:%s",
				question.Name, dns.TypeToString[question.Qtype])
			if err := responseDNSMsg(s.UDPConn, addr, msgCache, d.Address()); err != nil {
				log.Errorf("response dns msg err:%s", err.Error())
				return err
			}
			if strings.TrimSuffix(question.Name, ".") != ss.Server() {
				log.Debugf("renew dns cache for domain:%s", question.Name)
				ss.RenewDNSCache(question.Name, dns.TypeToString[question.Qtype], isDirect)
			}
			return nil
		}

		if isDirect {
			return ss.directUDPRelay(s, addr, d, true)
		}

		log.Debugf("rewrite dns dst addr to %s", DefaultDNSServer)
		rewrittenDst = DefaultDNSServer
	} else if err := ss.validateAddr(dst); err != nil {
		log.Warnf("validate udp dst:%v err:%v, data:%s", dst, err, string(d.Data))
		return errors.New("dst addr is invalid")
	}

	dstHost, _, _ := net.SplitHostPort(dst)
	if ss.HostAtCN(dstHost) && ss.Tun2socksStatusAuto() {
		return ss.directUDPRelay(s, addr, d, false)
	}

	var ch chan byte
	var hasAssoc bool

	portStr := strconv.FormatInt(int64(addr.Port), 10)
	asCh, ok := s.AssociatedUDP.Get(portStr)
	if ok {
		hasAssoc = true
		ch = asCh.(chan byte)
		log.Debugf("found the associate with tcp, src:%s, dst:%s", addr.String(), d.Address())
	} else {
		log.Debugf("the udp addr:%v doesn't associate with tcp, dst addr:%v", addr.String(), d.Address())
	}

	send := func(ue *UDPExchange, data []byte) error {
		select {
		case <-ch:
			return fmt.Errorf("the tcp that udp address %s associated closed", ue.ClientAddr.String())
		default:
		}
		_, err := ue.RemoteConn.Write(data)
		if err != nil {
			return err
		}
		log.Debugf("Sent UDP data to remote. client: %s", ue.ClientAddr.String())
		return nil
	}

	var ue *UDPExchange
	var src = addr.String()
	iue, ok := s.UDPExchanges.Get(src + dst)
	if ok {
		ue = iue.(*UDPExchange)
		return send(ue, d.Data)
	}

	pool := ss.Pool()
	if pool == nil {
		log.Errorf("failed to get pool, easyss is closed")
		return errors.New("easyss is closed")
	}

	stream, err := pool.Get()
	if err != nil {
		log.Errorf("get stream from pool failed:%+v", err)
		return err
	}

	if err := ss.handShakeWithRemote(stream, rewrittenDst, "udp"); err != nil {
		log.Errorf("handshake with remote server err:%v", err)
		if pc, ok := stream.(*easypool.PoolConn); ok {
			log.Debugf("mark pool conn stream unusable")
			pc.MarkUnusable()
			stream.Close()
		}
		return err
	}

	csStream, err := cipherstream.New(stream, ss.Password(), ss.Method(), "udp")
	if err != nil {
		log.Errorf("new cipherstream err:%+v, password:%v, method:%v",
			err, ss.Password(), ss.Method())
		return err
	}

	ue = &UDPExchange{
		ClientAddr: addr,
		RemoteConn: csStream,
	}
	if err := send(ue, d.Data); err != nil {
		markCipherStreamUnusable(ue.RemoteConn)
		ue.RemoteConn.Close()
		return err
	}
	s.UDPExchanges.Set(src+dst, ue, -1)

	var monitorCh = make(chan bool, 1)
	// monitor the assoc tcp connection to be closed and try to reuse the underlying connection
	go func() {
		var tryReuse bool
		select {
		case <-ch:
			expireConn(ue.RemoteConn)
			tryReuse = <-monitorCh
		case tryReuse = <-monitorCh:
		}

		if tryReuse {
			log.Debugf("udp request is finished, try reusing underlying tcp connection")
			buf := connStateBytes.Get(32)
			defer connStateBytes.Put(buf)

			state := NewConnState(FIN_WAIT1, buf)
			setCipherDeadline(ue.RemoteConn, ss.Timeout())
			for stateFn := state.fn; stateFn != nil; {
				stateFn = stateFn(ue.RemoteConn).fn
			}
			if state.err != nil {
				log.Infof("state err:%v, state:%v", state.err, state.state)
				markCipherStreamUnusable(ue.RemoteConn)
			} else {
				log.Debugf("underlying connection is health, so reuse it")
			}
		} else {
			markCipherStreamUnusable(ue.RemoteConn)
		}
		ue.RemoteConn.Close()
	}()

	go func(ue *UDPExchange, dst string) {
		defer func() {
			s.UDPExchanges.Delete(src + dst)
		}()

		var b = udpDataBytes.Get(MaxUDPDataSize)
		defer udpDataBytes.Put(b)
		for {
			select {
			case <-ch:
				log.Infof("the tcp that udp address %s associated closed", ue.ClientAddr.String())
				monitorCh <- true
				return
			default:
			}

			if !hasAssoc {
				if err := ue.RemoteConn.SetDeadline(time.Now().Add(10 * time.Second)); err != nil {
					log.Errorf("set the deadline for remote conn err:%v", err)
					monitorCh <- false
					return
				}
			}
			n, err := ue.RemoteConn.Read(b[:])
			if err != nil {
				if err == cipherstream.ErrTimeout {
					monitorCh <- true
				} else {
					monitorCh <- false
				}
				return
			}
			log.Debugf("got UDP data from remote. client: %v, data-len: %v", ue.ClientAddr.String(), len(b[0:n]))

			// if is dns response, set result to dns cache
			ss.SetDNSCacheIfNeeded(b[0:n], false)

			a, addr, port, err := socks5.ParseAddress(dst)
			if err != nil {
				monitorCh <- true
				log.Errorf("parse dst address err:%v", err)
				return
			}
			d1 := socks5.NewDatagram(a, addr, port, b[0:n])
			if _, err := s.UDPConn.WriteToUDP(d1.Bytes(), ue.ClientAddr); err != nil {
				monitorCh <- true
				return
			}
		}
	}(ue, dst)

	return nil
}

func (ss *Easyss) SetDNSCacheIfNeeded(udpResp []byte, isDirect bool) {
	msg := &dns.Msg{}
	if err := msg.Unpack(udpResp); err == nil && isDNSResponse(msg) {
		if err := ss.SetDNSCache(msg, false, isDirect); err != nil {
			log.Warnf("set dns cache err:%s", err.Error())
		} else {
			log.Debugf("set dns cache success for domain:%s, qtype:%s, isDirect:%v",
				msg.Question[0].Name, dns.TypeToString[msg.Question[0].Qtype], isDirect)
		}
	}
}

func responseDNSMsg(conn *net.UDPConn, localAddr *net.UDPAddr, msg *dns.Msg, remoteAddr string) error {
	a, _addr, port, _ := socks5.ParseAddress(remoteAddr)
	pack, _ := msg.Pack()
	d1 := socks5.NewDatagram(a, _addr, port, pack)

	_, err := conn.WriteToUDP(d1.Bytes(), localAddr)
	return err
}

func isDNSRequest(msg *dns.Msg) bool {
	if len(msg.Question) > 0 {
		return true
	}
	return false
}

func isDNSResponse(msg *dns.Msg) bool {
	if !msg.MsgHdr.Response || !isDNSRequest(msg) || len(msg.Answer) == 0 {
		return false
	}
	return true
}
<commit_msg>local_udp: fix dns request<commit_after>package easyss

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/miekg/dns"
	"github.com/nange/easypool"
	"github.com/nange/easyss/cipherstream"
	"github.com/nange/easyss/util"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"github.com/txthinking/socks5"
)

const (
	MaxUDPDataSize   = 65507
	DefaultDNSServer = "8.8.8.8:53"
)

var udpDataBytes = util.NewBytes(MaxUDPDataSize)

// UDPExchange used to store client address and remote connection
type UDPExchange struct {
	ClientAddr *net.UDPAddr
	RemoteConn net.Conn
}

func (ss *Easyss) UDPHandle(s *socks5.Server, addr *net.UDPAddr, d *socks5.Datagram) error {
	log.Debugf("enter udp handdle, local_addr:%v, remote_addr:%v", addr.String(), d.Address())

	dst := d.Address()
	rewrittenDst := dst

	msg := &dns.Msg{}
	err := msg.Unpack(d.Data)
	if err == nil && isDNSRequest(msg) {
		log.Infof("the udp request is dns proto, domain:%s, qtype:%s",
			msg.Question[0].Name, dns.TypeToString[msg.Question[0].Qtype])

		question := msg.Question[0]
		hostAtCN := ss.HostAtCN(strings.TrimSuffix(question.Name, "."))
		isDirect := hostAtCN && ss.Tun2socksStatusAuto()

		// find from dns cache first
		msgCache := ss.DNSCache(question.Name, dns.TypeToString[question.Qtype], isDirect)
		if msgCache != nil {
			msgCache.MsgHdr.Id = msg.MsgHdr.Id
			log.Infof("find msg from dns cache, write back directly, domain:%s, qtype:%s",
				question.Name, dns.TypeToString[question.Qtype])
			if err := responseDNSMsg(s.UDPConn, addr, msgCache, d.Address()); err != nil {
				log.Errorf("response dns msg err:%s", err.Error())
				return err
			}
			if strings.TrimSuffix(question.Name, ".") != ss.Server() {
				log.Debugf("renew dns cache for domain:%s", question.Name)
				ss.RenewDNSCache(question.Name, dns.TypeToString[question.Qtype], isDirect)
			}
			return nil
		}

		if isDirect {
			return ss.directUDPRelay(s, addr, d, true)
		}

		log.Debugf("rewrite dns dst addr to %s", DefaultDNSServer)
		rewrittenDst = DefaultDNSServer
	} else if err := ss.validateAddr(dst); err != nil {
		log.Warnf("validate udp dst:%v err:%v, data:%s", dst, err, string(d.Data))
		return errors.New("dst addr is invalid")
	}

	dstHost, _, _ := net.SplitHostPort(rewrittenDst)
	if ss.HostAtCN(dstHost) && ss.Tun2socksStatusAuto() {
		return ss.directUDPRelay(s, addr, d, isDNSRequest(msg))
	}

	var ch chan byte
	var hasAssoc bool

	portStr := strconv.FormatInt(int64(addr.Port), 10)
	asCh, ok := s.AssociatedUDP.Get(portStr)
	if ok {
		hasAssoc = true
		ch = asCh.(chan byte)
		log.Debugf("found the associate with tcp, src:%s, dst:%s", addr.String(), d.Address())
	} else {
		log.Debugf("the udp addr:%v doesn't associate with tcp, dst addr:%v", addr.String(), d.Address())
	}

	send := func(ue *UDPExchange, data []byte) error {
		select {
		case <-ch:
			return fmt.Errorf("the tcp that udp address %s associated closed", ue.ClientAddr.String())
		default:
		}
		_, err := ue.RemoteConn.Write(data)
		if err != nil {
			return err
		}
		log.Debugf("Sent UDP data to remote. client: %s", ue.ClientAddr.String())
		return nil
	}

	var ue *UDPExchange
	var src = addr.String()
	iue, ok := s.UDPExchanges.Get(src + dst)
	if ok {
		ue = iue.(*UDPExchange)
		return send(ue, d.Data)
	}

	pool := ss.Pool()
	if pool == nil {
		log.Errorf("failed to get pool, easyss is closed")
		return errors.New("easyss is closed")
	}

	stream, err := pool.Get()
	if err != nil {
		log.Errorf("get stream from pool failed:%+v", err)
		return err
	}

	if err := ss.handShakeWithRemote(stream, rewrittenDst, "udp"); err != nil {
		log.Errorf("handshake with remote server err:%v", err)
		if pc, ok := stream.(*easypool.PoolConn); ok {
			log.Debugf("mark pool conn stream unusable")
			pc.MarkUnusable()
			stream.Close()
		}
		return err
	}

	csStream, err := cipherstream.New(stream, ss.Password(), ss.Method(), "udp")
	if err != nil {
		log.Errorf("new cipherstream err:%+v, password:%v, method:%v",
			err, ss.Password(), ss.Method())
		return err
	}

	ue = &UDPExchange{
		ClientAddr: addr,
		RemoteConn: csStream,
	}
	if err := send(ue, d.Data); err != nil {
		markCipherStreamUnusable(ue.RemoteConn)
		ue.RemoteConn.Close()
		return err
	}
	s.UDPExchanges.Set(src+dst, ue, -1)

	var monitorCh = make(chan bool, 1)
	// monitor the assoc tcp connection to be closed and try to reuse the underlying connection
	go func() {
		var tryReuse bool
		select {
		case <-ch:
			expireConn(ue.RemoteConn)
			tryReuse = <-monitorCh
		case tryReuse = <-monitorCh:
		}

		if tryReuse {
			log.Debugf("udp request is finished, try reusing underlying tcp connection")
			buf := connStateBytes.Get(32)
			defer connStateBytes.Put(buf)

			state := NewConnState(FIN_WAIT1, buf)
			setCipherDeadline(ue.RemoteConn, ss.Timeout())
			for stateFn := state.fn; stateFn != nil; {
				stateFn = stateFn(ue.RemoteConn).fn
			}
			if state.err != nil {
				log.Infof("state err:%v, state:%v", state.err, state.state)
				markCipherStreamUnusable(ue.RemoteConn)
			} else {
				log.Debugf("underlying connection is health, so reuse it")
			}
		} else {
			markCipherStreamUnusable(ue.RemoteConn)
		}
		ue.RemoteConn.Close()
	}()

	go func(ue *UDPExchange, dst string) {
		defer func() {
			s.UDPExchanges.Delete(src + dst)
		}()

		var b = udpDataBytes.Get(MaxUDPDataSize)
		defer udpDataBytes.Put(b)
		for {
			select {
			case <-ch:
				log.Infof("the tcp that udp address %s associated closed", ue.ClientAddr.String())
				monitorCh <- true
				return
			default:
			}

			if !hasAssoc {
				if err := ue.RemoteConn.SetDeadline(time.Now().Add(10 * time.Second)); err != nil {
					log.Errorf("set the deadline for remote conn err:%v", err)
					monitorCh <- false
					return
				}
			}
			n, err := ue.RemoteConn.Read(b[:])
			if err != nil {
				if err == cipherstream.ErrTimeout {
					monitorCh <- true
				} else {
					monitorCh <- false
				}
				return
			}
			log.Debugf("got UDP data from remote. client: %v, data-len: %v", ue.ClientAddr.String(), len(b[0:n]))

			// if is dns response, set result to dns cache
			ss.SetDNSCacheIfNeeded(b[0:n], false)

			a, addr, port, err := socks5.ParseAddress(dst)
			if err != nil {
				monitorCh <- true
				log.Errorf("parse dst address err:%v", err)
				return
			}
			d1 := socks5.NewDatagram(a, addr, port, b[0:n])
			if _, err := s.UDPConn.WriteToUDP(d1.Bytes(), ue.ClientAddr); err != nil {
				monitorCh <- true
				return
			}
		}
	}(ue, dst)

	return nil
}

func (ss *Easyss) SetDNSCacheIfNeeded(udpResp []byte, isDirect bool) {
	msg := &dns.Msg{}
	if err := msg.Unpack(udpResp); err == nil && isDNSResponse(msg) {
		if err := ss.SetDNSCache(msg, false, isDirect); err != nil {
			log.Warnf("set dns cache err:%s", err.Error())
		} else {
			log.Debugf("set dns cache success for domain:%s, qtype:%s, isDirect:%v",
				msg.Question[0].Name, dns.TypeToString[msg.Question[0].Qtype], isDirect)
		}
	}
}

func responseDNSMsg(conn *net.UDPConn, localAddr *net.UDPAddr, msg *dns.Msg, remoteAddr string) error {
	a, _addr, port, _ := socks5.ParseAddress(remoteAddr)
	pack, _ := msg.Pack()
	d1 := socks5.NewDatagram(a, _addr, port, pack)

	_, err := conn.WriteToUDP(d1.Bytes(), localAddr)
	return err
}

func isDNSRequest(msg *dns.Msg) bool {
	if msg == nil {
		return false
	}
	if len(msg.Question) > 0 {
		return true
	}
	return false
}

func isDNSResponse(msg *dns.Msg) bool {
	if msg == nil {
		return false
	}
	if !msg.MsgHdr.Response || !isDNSRequest(msg) || len(msg.Answer) == 0 {
		return false
	}
	return true
}
<|endoftext|>