<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 fuseops

import (
	"flag"
	"fmt"
	"reflect"
	"strings"
	"sync"
	"time"

	"github.com/jacobsa/bazilfuse"
	"github.com/jacobsa/reqtrace"
	"golang.org/x/net/context"
	"golang.org/x/sys/unix"
)

var fTraceByPID = flag.Bool(
	"fuse.trace_by_pid",
	false,
	"Enable a hacky mode that uses reqtrace to group all ops from each "+
		"individual PID. Not a good idea to use in production; races, bugs, and "+
		"resource leaks likely lurk.")

// A helper for embedding common behavior.
type commonOp struct {
	opType      string
	r           bazilfuse.Request
	log         func(int, string, ...interface{})
	opsInFlight *sync.WaitGroup

	ctx    context.Context
	report reqtrace.ReportFunc
}

func describeOpType(t reflect.Type) (desc string) {
	name := t.String()

	// The usual case: a string that looks like "*fuseops.GetInodeAttributesOp".
	const prefix = "*fuseops."
	const suffix = "Op"
	if strings.HasPrefix(name, prefix) && strings.HasSuffix(name, suffix) {
		desc = name[len(prefix) : len(name)-len(suffix)]
		return
	}

	// Otherwise, it's not clear what to do.
	desc = t.String()

	return
}

var gPIDMapMu sync.Mutex

// A map from PID to a traced context for that PID.
//
// GUARDED_BY(gPIDMapMu)
var gPIDMap = make(map[int]context.Context)

// Wait until the process completes, then close off the trace and remove the
// context from the map.
func reportWhenPIDGone(
	pid int,
	ctx context.Context,
	report reqtrace.ReportFunc) {
	// HACK(jacobsa): Poll for completion.
	const pollPeriod = 50 * time.Millisecond
	for {
		err := unix.Kill(pid, 0)
		if err == unix.ESRCH {
			break
		}

		if err != nil {
			panic(fmt.Errorf("Kill(%v): %v", pid, err))
		}

		time.Sleep(pollPeriod)
	}

	// Finish up.
	report(nil)

	gPIDMapMu.Lock()
	delete(gPIDMap, pid)
	gPIDMapMu.Unlock()
}

func (o *commonOp) maybeTraceByPID(
	in context.Context,
	pid int) (out context.Context) {
	// Is there anything to do?
	if !reqtrace.Enabled() || !*fTraceByPID {
		out = in
		return
	}

	gPIDMapMu.Lock()
	defer gPIDMapMu.Unlock()

	// Do we already have a traced context for this PID?
	if existing, ok := gPIDMap[pid]; ok {
		out = existing
		return
	}

	// Set up a new one and stick it in the map.
	var report reqtrace.ReportFunc
	out, report = reqtrace.Trace(in, fmt.Sprintf("PID %v", pid))
	gPIDMap[pid] = out

	// Ensure we close the trace and remove it from the map eventually.
	go reportWhenPIDGone(pid, out, report)

	return
}

func (o *commonOp) init(
	ctx context.Context,
	opType reflect.Type,
	r bazilfuse.Request,
	log func(int, string, ...interface{}),
	opsInFlight *sync.WaitGroup) {
	// Set up a context that reflects per-PID tracing if appropriate.
	ctx = o.maybeTraceByPID(ctx, int(r.Hdr().Pid))

	// Initialize basic fields.
	o.opType = describeOpType(opType)
	o.r = r
	o.log = log
	o.opsInFlight = opsInFlight

	// Set up a trace span for this op.
	o.ctx, o.report = reqtrace.StartSpan(ctx, o.opType)
}

func (o *commonOp) Header() OpHeader {
	bh := o.r.Hdr()
	return OpHeader{
		Uid: bh.Uid,
		Gid: bh.Gid,
	}
}

func (o *commonOp) Context() context.Context {
	return o.ctx
}

func (o *commonOp) Logf(format string, v ...interface{}) {
	const calldepth = 2
	o.log(calldepth, format, v...)
}

func (o *commonOp) respondErr(err error) {
	if err == nil {
		panic("Expect non-nil here.")
	}

	o.report(err)

	o.Logf(
		"-> (%s) error: %v",
		o.opType,
		err)

	o.r.RespondError(err)
}

// Respond with the supplied response struct, which must be accepted by a
// method called Respond on o.r.
//
// Special case: nil means o.r.Respond accepts no parameters.
func (o *commonOp) respond(resp interface{}) {
	// We were successful.
	o.report(nil)

	// Find the Respond method.
	v := reflect.ValueOf(o.r)
	respond := v.MethodByName("Respond")

	// Special case: handle successful ops with no response struct.
	if resp == nil {
		o.Logf("-> (%s) OK", o.opType)
		respond.Call([]reflect.Value{})
		return
	}

	// Otherwise, pass along the response struct.
	o.Logf("-> %v", resp)
	respond.Call([]reflect.Value{reflect.ValueOf(resp)})
}
<commit_msg>Handle permission errors.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 fuseops

import (
	"flag"
	"fmt"
	"log"
	"reflect"
	"strings"
	"sync"
	"time"

	"github.com/jacobsa/bazilfuse"
	"github.com/jacobsa/reqtrace"
	"golang.org/x/net/context"
	"golang.org/x/sys/unix"
)

var fTraceByPID = flag.Bool(
	"fuse.trace_by_pid",
	false,
	"Enable a hacky mode that uses reqtrace to group all ops from each "+
		"individual PID. Not a good idea to use in production; races, bugs, and "+
		"resource leaks likely lurk.")

// A helper for embedding common behavior.
type commonOp struct {
	opType      string
	r           bazilfuse.Request
	log         func(int, string, ...interface{})
	opsInFlight *sync.WaitGroup

	ctx    context.Context
	report reqtrace.ReportFunc
}

func describeOpType(t reflect.Type) (desc string) {
	name := t.String()

	// The usual case: a string that looks like "*fuseops.GetInodeAttributesOp".
	const prefix = "*fuseops."
	const suffix = "Op"
	if strings.HasPrefix(name, prefix) && strings.HasSuffix(name, suffix) {
		desc = name[len(prefix) : len(name)-len(suffix)]
		return
	}

	// Otherwise, it's not clear what to do.
	desc = t.String()

	return
}

var gPIDMapMu sync.Mutex

// A map from PID to a traced context for that PID.
//
// GUARDED_BY(gPIDMapMu)
var gPIDMap = make(map[int]context.Context)

// Wait until the process completes, then close off the trace and remove the
// context from the map.
func reportWhenPIDGone(
	pid int,
	ctx context.Context,
	report reqtrace.ReportFunc) {
	// HACK(jacobsa): Poll until the process no longer exists.
	const pollPeriod = 50 * time.Millisecond
	for {
		// The man page for kill(2) says that if the signal is zero, then "no
		// signal is sent, but error checking is still performed; this can be used
		// to check for the existence of a process ID".
		err := unix.Kill(pid, 0)

		// ESRCH means the process is gone.
		if err == unix.ESRCH {
			break
		}

		// If we receive EPERM, we're not going to be able to do what we want. We
		// don't really have any choice but to print info and leak.
		if err == unix.EPERM {
			log.Printf("Failed to kill(2) PID %v; no permissions. Leaking trace.", pid)
			return
		}

		// Otherwise, panic.
		if err != nil {
			panic(fmt.Errorf("Kill(%v): %v", pid, err))
		}

		time.Sleep(pollPeriod)
	}

	// Finish up.
	report(nil)

	gPIDMapMu.Lock()
	delete(gPIDMap, pid)
	gPIDMapMu.Unlock()
}

func (o *commonOp) maybeTraceByPID(
	in context.Context,
	pid int) (out context.Context) {
	// Is there anything to do?
	if !reqtrace.Enabled() || !*fTraceByPID {
		out = in
		return
	}

	gPIDMapMu.Lock()
	defer gPIDMapMu.Unlock()

	// Do we already have a traced context for this PID?
	if existing, ok := gPIDMap[pid]; ok {
		out = existing
		return
	}

	// Set up a new one and stick it in the map.
	var report reqtrace.ReportFunc
	out, report = reqtrace.Trace(in, fmt.Sprintf("PID %v", pid))
	gPIDMap[pid] = out

	// Ensure we close the trace and remove it from the map eventually.
	go reportWhenPIDGone(pid, out, report)

	return
}

func (o *commonOp) init(
	ctx context.Context,
	opType reflect.Type,
	r bazilfuse.Request,
	log func(int, string, ...interface{}),
	opsInFlight *sync.WaitGroup) {
	// Set up a context that reflects per-PID tracing if appropriate.
	ctx = o.maybeTraceByPID(ctx, int(r.Hdr().Pid))

	// Initialize basic fields.
	o.opType = describeOpType(opType)
	o.r = r
	o.log = log
	o.opsInFlight = opsInFlight

	// Set up a trace span for this op.
	o.ctx, o.report = reqtrace.StartSpan(ctx, o.opType)
}

func (o *commonOp) Header() OpHeader {
	bh := o.r.Hdr()
	return OpHeader{
		Uid: bh.Uid,
		Gid: bh.Gid,
	}
}

func (o *commonOp) Context() context.Context {
	return o.ctx
}

func (o *commonOp) Logf(format string, v ...interface{}) {
	const calldepth = 2
	o.log(calldepth, format, v...)
}

func (o *commonOp) respondErr(err error) {
	if err == nil {
		panic("Expect non-nil here.")
	}

	o.report(err)

	o.Logf(
		"-> (%s) error: %v",
		o.opType,
		err)

	o.r.RespondError(err)
}

// Respond with the supplied response struct, which must be accepted by a
// method called Respond on o.r.
//
// Special case: nil means o.r.Respond accepts no parameters.
func (o *commonOp) respond(resp interface{}) {
	// We were successful.
	o.report(nil)

	// Find the Respond method.
	v := reflect.ValueOf(o.r)
	respond := v.MethodByName("Respond")

	// Special case: handle successful ops with no response struct.
	if resp == nil {
		o.Logf("-> (%s) OK", o.opType)
		respond.Call([]reflect.Value{})
		return
	}

	// Otherwise, pass along the response struct.
	o.Logf("-> %v", resp)
	respond.Call([]reflect.Value{reflect.ValueOf(resp)})
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/fcgi"
	"net/url"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
)

var (
	form          url.Values
	header        http.Header
	cookies       []*http.Cookie
	writer        http.ResponseWriter
	request       http.Request
	exit_error    bool
	server        *GochanServer
	referrerRegex *regexp.Regexp
)

type GochanServer struct {
	writer     http.ResponseWriter
	request    http.Request
	namespaces map[string]func(http.ResponseWriter, *http.Request, interface{})
}

func (s GochanServer) AddNamespace(base_path string, namespace_function func(http.ResponseWriter, *http.Request, interface{})) {
	s.namespaces[base_path] = namespace_function
}

func (s GochanServer) getFileData(writer http.ResponseWriter, url string) ([]byte, bool) {
	var file_bytes []byte
	filepath := path.Join(config.DocumentRoot, url)
	results, err := os.Stat(filepath)
	if err != nil {
		// the requested path isn't a file or directory, 404
		return file_bytes, false
	} else {
		//the file exists, or there is a folder here
		if results.IsDir() {
			newpath := ""

			//check to see if one of the specified index pages exists
			for _, value := range(config.FirstPage) {
				newpath = path.Join(filepath, value)
				_, err := os.Stat(newpath)
				if err == nil {
					// serve the index page
					writer.Header().Add("Cache-Control", "max-age=5, must-revalidate")
					file_bytes, err = ioutil.ReadFile(newpath)
					return file_bytes, true
				}
			}
		} else {
			//the file exists, and is not a folder
			file_bytes, err = ioutil.ReadFile(filepath)
			extension := getFileExtension(url)
			switch {
			case extension == "png":
				writer.Header().Add("Content-Type", "image/png")
				writer.Header().Add("Cache-Control", "max-age=86400")
			case extension == "gif":
				writer.Header().Add("Content-Type", "image/gif")
				writer.Header().Add("Cache-Control", "max-age=86400")
			case extension == "jpg":
				writer.Header().Add("Content-Type", "image/jpeg")
				writer.Header().Add("Cache-Control", "max-age=86400")
			case extension == "css":
				writer.Header().Add("Content-Type", "text/css")
				writer.Header().Add("Cache-Control", "max-age=43200")
			case extension == "js":
				writer.Header().Add("Content-Type", "text/javascript")
				writer.Header().Add("Cache-Control", "max-age=43200")
			case extension == "json":
				writer.Header().Add("Content-Type", "application/json")
				writer.Header().Add("Cache-Control", "max-age=5, must-revalidate")
			}
			if extension == "html" || extension == "htm" {
				writer.Header().Add("Cache-Control", "max-age=5, must-revalidate")
			}
			access_log.Print("Success: 200 from " + getRealIP(&request) + " @ " + request.RequestURI)
			return file_bytes, true
		}
	}
	return file_bytes, false
}

func (s GochanServer) Redirect(location string) {
	http.Redirect(writer, &request, location, http.StatusFound)
}

func (s GochanServer) serve404(writer http.ResponseWriter, request *http.Request) {
	error_page, err := ioutil.ReadFile(config.DocumentRoot + "/error/404.html")
	if err != nil {
		writer.Write([]byte("Requested page not found, and 404 error page not found"))
	} else {
		writer.Write(error_page)
	}
	error_log.Print("Error: 404 Not Found from " + getRealIP(request) + " @ " + request.RequestURI)
}

func (s GochanServer) ServeErrorPage(writer http.ResponseWriter, err string) {
	error_page_bytes, _ := ioutil.ReadFile("templates/error.html")
	error_page := string(error_page_bytes)
	error_page = strings.Replace(error_page, "{ERRORTEXT}", err, -1)
	writer.Write([]byte(error_page))
	exit_error = true
}

func (s GochanServer) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	for name, namespace_function := range s.namespaces {
		//if len(request.URL)
		if request.URL.Path == "/"+name {
			namespace_function(writer, request, nil)
			return
		}
	}
	fb, found := s.getFileData(writer, request.URL.Path)
	writer.Header().Add("Cache-Control", "max-age=86400")
	if !found {
		s.serve404(writer, request)
		return
	}
	writer.Write(fb)
}

func initServer() {
	listener, err := net.Listen("tcp", config.ListenIP+":"+strconv.Itoa(config.Port))
	if err != nil {
		fmt.Printf("Failed listening on %s:%d, see log for details", config.ListenIP, config.Port)
		error_log.Fatal(err.Error())
	}
	server = new(GochanServer)
	server.namespaces = make(map[string]func(http.ResponseWriter, *http.Request, interface{}))

	// Check if Akismet API key is usable at startup.
	if config.AkismetAPIKey != "" {
		checkAkismetAPIKey()
	}

	// Compile regex for checking referrers.
	referrerRegex = regexp.MustCompile(config.DomainRegex)

	testfunc := func(writer http.ResponseWriter, response *http.Request, data interface{}) {
		if writer != nil {
			writer.Write([]byte("hahahaha"))
		}
	}
	server.AddNamespace("example", testfunc)
	server.AddNamespace("manage", callManageFunction)
	server.AddNamespace("post", makePost)
	server.AddNamespace("util", utilHandler)
	// eventually plugins will be able to register new namespaces. Or they will be restricted to something like /plugin
	if config.UseFastCGI {
		fcgi.Serve(listener, server)
	} else {
		http.Serve(listener, server)
	}
}

func getRealIP(r *http.Request) (ip string) {
	// HTTP_CF_CONNECTING_IP > X-Forwarded-For > RemoteAddr
	if r.Header.Get("HTTP_CF_CONNECTING_IP") != "" {
		ip = r.Header.Get("HTTP_CF_CONNECTING_IP")
	} else {
		if r.Header.Get("X-Forwarded-For") != "" {
			ip = r.Header.Get("X-Forwarded-For")
		} else {
			ip = r.RemoteAddr
		}
	}
	return
}

func validReferrer(request http.Request) (valid bool) {

	referrerCopy := request.Referer()
	valid = referrerRegex.MatchString(referrerCopy)
	// Old Referrer check.
	// valid = !(request.Referer() == "" || len(request.Referer()) < len(config.SiteDomain) || request.Referer()[7:len(config.SiteDomain)+7] != config.SiteDomain)
	return
}

// register /util handler
func utilHandler(writer http.ResponseWriter, request *http.Request, data interface{}) {
	writer.Header().Add("Content-Type", "text/css")
	action := request.FormValue("action")
	board := request.FormValue("board")
	var err error
	if action == "" && request.PostFormValue("delete_btn") != "Delete" && request.PostFormValue("report_btn") != "Report" {
		http.Redirect(writer, request, path.Join(config.SiteWebfolder, "/"), http.StatusFound)
		return
	}
	var posts_arr []string
	for key, _ := range request.PostForm {
		if strings.Index(key, "check") == 0 {
			posts_arr = append(posts_arr, key[5:])
		}
	}
	if request.PostFormValue("delete_btn") == "Delete" {
		// Delete a post or thread
		file_only := request.FormValue("fileonly") == "on"
		password := md5_sum(request.FormValue("password"))
		rank := getStaffRank()

		if request.FormValue("password") == "" && rank == 0 {
			server.ServeErrorPage(writer, "Password required for post deletion")
			return
		}

		for _, post := range posts_arr {
			var parent_id int
			var filename string
			var filetype string
			var password_checksum string
			var board_id int
			//post_int,err := strconv.Atoi(post)

			err = db.QueryRow("SELECT `parentid`,`filename`,`password` FROM `"+config.DBprefix+"posts` WHERE `id` = "+post+" AND `deleted_timestamp` = '"+nil_timestamp+"'").Scan(&parent_id, &filename, &password_checksum)
			if err == sql.ErrNoRows {
				//the post has already been deleted
				writer.Header().Add("refresh", "3;url="+request.Referer())
				fmt.Fprintf(writer, "%s has already been deleted or is a post in a deleted thread.\n
", post)
				continue
			}
			if err != nil {
				server.ServeErrorPage(writer, err.Error())
				return
			}

			err = db.QueryRow("SELECT `id` FROM `" + config.DBprefix + "boards` WHERE `dir` = '" + board + "'").Scan(&board_id)
			if err != nil {
				server.ServeErrorPage(writer, err.Error())
				return
			}

			if password != password_checksum && rank == 0 {
				fmt.Fprintf(writer, "Incorrect password for %s\n", post)
				continue
			}

			if file_only {

				if filename != "" && filename != "deleted" {
					filetype = filename[strings.Index(filename, ".")+1:]
					filename = filename[:strings.Index(filename, ".")]
					err := os.Remove(path.Join(config.DocumentRoot, board, "/src/"+filename+"."+filetype))
					if err != nil {
						server.ServeErrorPage(writer, err.Error())
						return
					}
					err = os.Remove(path.Join(config.DocumentRoot, board, "/thumb/"+filename+"t."+filetype))
					if err != nil {
						server.ServeErrorPage(writer, err.Error())
						return
					}
					_, err = db.Exec("UPDATE `" + config.DBprefix + "posts` SET `filename` = 'deleted' WHERE `id` = " + post)
					if err != nil {
						server.ServeErrorPage(writer, err.Error())
						return
					}
				}
				writer.Header().Add("refresh", "3;url="+request.Referer())
				fmt.Fprintf(writer, "Attached image from %s deleted successfully
\n<meta http-equiv=\"refresh\" content=\"1;url="+config.DocumentRoot+"/"+board+"/\">", post)
			} else {

				// delete the post
				_, err = db.Exec("UPDATE `" + config.DBprefix + "posts` SET `deleted_timestamp` = '" + getSQLDateTime() + "' WHERE `id` = " + post)
				if parent_id == 0 {
					err = os.Remove(path.Join(config.DocumentRoot, board, "/res/"+post+".html"))
				} else {
					_board, _ := getBoardArr("`id` = " + strconv.Itoa(board_id))
					buildBoardPages(&_board[0])
				}

				// if the deleted post is actually a thread, delete its posts
				_, _ = db.Exec("UPDATE `" + config.DBprefix + "posts` SET `deleted_timestamp` = '" + getSQLDateTime() + "' WHERE `parentid` = " + post)
				if err != nil {
					server.ServeErrorPage(writer, err.Error())
					return
				}

				// delete the
				var deleted_filename string
				err = db.QueryRow("SELECT `filename` FROM `" + config.DBprefix + "posts` WHERE `id` = " + post + " AND `filename` != ''").Scan(&deleted_filename)
				if err == nil {
					os.Remove(path.Join(config.DocumentRoot, board, "/src/", deleted_filename))
					os.Remove(path.Join(config.DocumentRoot, board, "/thumb/", strings.Replace(deleted_filename, ".", "t.", -1)))
					os.Remove(path.Join(config.DocumentRoot, board, "/thumb/", strings.Replace(deleted_filename, ".", "c.", -1)))
				}

				err = db.QueryRow("SELECT `filename` FROM `" + config.DBprefix + "posts` WHERE `parentid` = " + post + " AND `filename` != ''").Scan(&deleted_filename)
				if err == nil {
					os.Remove(path.Join(config.DocumentRoot, board, "/src/", deleted_filename))
					os.Remove(path.Join(config.DocumentRoot, board, "/thumb/", strings.Replace(deleted_filename, ".", "t.", -1)))
					os.Remove(path.Join(config.DocumentRoot, board, "/thumb/", strings.Replace(deleted_filename, ".", "c.", -1)))
				}

				buildBoards(false, board_id)

				writer.Header().Add("refresh", "3;url="+request.Referer())
				fmt.Fprintf(writer, "%s deleted successfully\n
", post)
			}
		}
	}
}
<commit_msg>Update server.go<commit_after>package main

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/fcgi"
	"net/url"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
)

var (
	form          url.Values
	header        http.Header
	cookies       []*http.Cookie
	writer        http.ResponseWriter
	request       http.Request
	exit_error    bool
	server        *GochanServer
	referrerRegex *regexp.Regexp
)

type GochanServer struct {
	writer     http.ResponseWriter
	request    http.Request
	namespaces map[string]func(http.ResponseWriter, *http.Request, interface{})
}

func (s GochanServer) AddNamespace(base_path string, namespace_function func(http.ResponseWriter, *http.Request, interface{})) {
	s.namespaces[base_path] = namespace_function
}

func (s GochanServer) getFileData(writer http.ResponseWriter, url string) ([]byte, bool) {
	var file_bytes []byte
	filepath := path.Join(config.DocumentRoot, url)
	results, err := os.Stat(filepath)
	if err != nil {
		// the requested path isn't a file or directory, 404
		return file_bytes, false
	} else {
		//the file exists, or there is a folder here
		if results.IsDir() {
			newpath := ""

			//check to see if one of the specified index pages exists
			for _, value := range(config.FirstPage) {
				newpath = path.Join(filepath, value)
				_, err := os.Stat(newpath)
				if err == nil {
					// serve the index page
					writer.Header().Add("Cache-Control", "max-age=5, must-revalidate")
					file_bytes, err = ioutil.ReadFile(newpath)
					return file_bytes, true
				}
			}
		} else {
			//the file exists, and is not a folder
			file_bytes, err = ioutil.ReadFile(filepath)
			extension := getFileExtension(url)
			switch extension {
			case "png":
				writer.Header().Add("Content-Type", "image/png")
				writer.Header().Add("Cache-Control", "max-age=86400")
			case "gif":
				writer.Header().Add("Content-Type", "image/gif")
				writer.Header().Add("Cache-Control", "max-age=86400")
			case "jpg":
				writer.Header().Add("Content-Type", "image/jpeg")
				writer.Header().Add("Cache-Control", "max-age=86400")
			case "css":
				writer.Header().Add("Content-Type", "text/css")
				writer.Header().Add("Cache-Control", "max-age=43200")
			case "js":
				writer.Header().Add("Content-Type", "text/javascript")
				writer.Header().Add("Cache-Control", "max-age=43200")
			case "json":
				writer.Header().Add("Content-Type", "application/json")
				writer.Header().Add("Cache-Control", "max-age=5, must-revalidate")
			}
			if strings.HasPrefix(extension, "htm") {
				writer.Header().Add("Cache-Control", "max-age=5, must-revalidate")
			}

			access_log.Print("Success: 200 from " + getRealIP(&request) + " @ " + request.RequestURI)
			return file_bytes, true
		}
	}
	return file_bytes, false
}

func (s GochanServer) Redirect(location string) {
	http.Redirect(writer, &request, location, http.StatusFound)
}

func (s GochanServer) serve404(writer http.ResponseWriter, request *http.Request) {
	error_page, err := ioutil.ReadFile(config.DocumentRoot + "/error/404.html")
	if err != nil {
		writer.Write([]byte("Requested page not found, and 404 error page not found"))
	} else {
		writer.Write(error_page)
	}
	error_log.Print("Error: 404 Not Found from " + getRealIP(request) + " @ " + request.RequestURI)
}

func (s GochanServer) ServeErrorPage(writer http.ResponseWriter, err string) {
	error_page_bytes, _ := ioutil.ReadFile("templates/error.html")
	error_page := string(error_page_bytes)
	error_page = strings.Replace(error_page, "{ERRORTEXT}", err, -1)
	writer.Write([]byte(error_page))
	exit_error = true
}

func (s GochanServer) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	for name, namespace_function := range s.namespaces {
		//if len(request.URL)
		if request.URL.Path == "/"+name {
			namespace_function(writer, request, nil)
			return
		}
	}
	fb, found := s.getFileData(writer, request.URL.Path)
	writer.Header().Add("Cache-Control", "max-age=86400")
	if !found {
		s.serve404(writer, request)
		return
	}
	writer.Write(fb)
}

func initServer() {
	listener, err := net.Listen("tcp", config.ListenIP+":"+strconv.Itoa(config.Port))
	if err != nil {
		fmt.Printf("Failed listening on %s:%d, see log for details", config.ListenIP, config.Port)
		error_log.Fatal(err.Error())
	}
	server = new(GochanServer)
	server.namespaces = make(map[string]func(http.ResponseWriter, *http.Request, interface{}))

	// Check if Akismet API key is usable at startup.
	if config.AkismetAPIKey != "" {
		checkAkismetAPIKey()
	}

	// Compile regex for checking referrers.
	referrerRegex = regexp.MustCompile(config.DomainRegex)

	testfunc := func(writer http.ResponseWriter, response *http.Request, data interface{}) {
		if writer != nil {
			writer.Write([]byte("hahahaha"))
		}
	}
	server.AddNamespace("example", testfunc)
	server.AddNamespace("manage", callManageFunction)
	server.AddNamespace("post", makePost)
	server.AddNamespace("util", utilHandler)
	// eventually plugins will be able to register new namespaces. Or they will be restricted to something like /plugin
	if config.UseFastCGI {
		fcgi.Serve(listener, server)
	} else {
		http.Serve(listener, server)
	}
}

func getRealIP(r *http.Request) (ip string) {
	// HTTP_CF_CONNECTING_IP > X-Forwarded-For > RemoteAddr
	if r.Header.Get("HTTP_CF_CONNECTING_IP") != "" {
		ip = r.Header.Get("HTTP_CF_CONNECTING_IP")
	} else {
		if r.Header.Get("X-Forwarded-For") != "" {
			ip = r.Header.Get("X-Forwarded-For")
		} else {
			ip = r.RemoteAddr
		}
	}
	return
}

func validReferrer(request http.Request) (valid bool) {

	referrerCopy := request.Referer()
	valid = referrerRegex.MatchString(referrerCopy)
	// Old Referrer check.
	// valid = !(request.Referer() == "" || len(request.Referer()) < len(config.SiteDomain) || request.Referer()[7:len(config.SiteDomain)+7] != config.SiteDomain)
	return
}

// register /util handler
func utilHandler(writer http.ResponseWriter, request *http.Request, data interface{}) {
	writer.Header().Add("Content-Type", "text/css")
	action := request.FormValue("action")
	board := request.FormValue("board")
	var err error
	if action == "" && request.PostFormValue("delete_btn") != "Delete" && request.PostFormValue("report_btn") != "Report" {
		http.Redirect(writer, request, path.Join(config.SiteWebfolder, "/"), http.StatusFound)
		return
	}
	var posts_arr []string
	for key, _ := range request.PostForm {
		if strings.Index(key, "check") == 0 {
			posts_arr = append(posts_arr, key[5:])
		}
	}
	if request.PostFormValue("delete_btn") == "Delete" {
		// Delete a post or thread
		file_only := request.FormValue("fileonly") == "on"
		password := md5_sum(request.FormValue("password"))
		rank := getStaffRank()

		if request.FormValue("password") == "" && rank == 0 {
			server.ServeErrorPage(writer, "Password required for post deletion")
			return
		}

		for _, post := range posts_arr {
			var parent_id int
			var filename string
			var filetype string
			var password_checksum string
			var board_id int
			//post_int,err := strconv.Atoi(post)

			err = db.QueryRow("SELECT `parentid`,`filename`,`password` FROM `"+config.DBprefix+"posts` WHERE `id` = "+post+" AND `deleted_timestamp` = '"+nil_timestamp+"'").Scan(&parent_id, &filename, &password_checksum)
			if err == sql.ErrNoRows {
				//the post has already been deleted
				writer.Header().Add("refresh", "3;url="+request.Referer())
				fmt.Fprintf(writer, "%s has already been deleted or is a post in a deleted thread.\n
", post)
				continue
			}
			if err != nil {
				server.ServeErrorPage(writer, err.Error())
				return
			}

			err = db.QueryRow("SELECT `id` FROM `" + config.DBprefix + "boards` WHERE `dir` = '" + board + "'").Scan(&board_id)
			if err != nil {
				server.ServeErrorPage(writer, err.Error())
				return
			}

			if password != password_checksum && rank == 0 {
				fmt.Fprintf(writer, "Incorrect password for %s\n", post)
				continue
			}

			if file_only {

				if filename != "" && filename != "deleted" {
					filetype = filename[strings.Index(filename, ".")+1:]
					filename = filename[:strings.Index(filename, ".")]
					err := os.Remove(path.Join(config.DocumentRoot, board, "/src/"+filename+"."+filetype))
					if err != nil {
						server.ServeErrorPage(writer, err.Error())
						return
					}
					err = os.Remove(path.Join(config.DocumentRoot, board, "/thumb/"+filename+"t."+filetype))
					if err != nil {
						server.ServeErrorPage(writer, err.Error())
						return
					}
					_, err = db.Exec("UPDATE `" + config.DBprefix + "posts` SET `filename` = 'deleted' WHERE `id` = " + post)
					if err != nil {
						server.ServeErrorPage(writer, err.Error())
						return
					}
				}
				writer.Header().Add("refresh", "3;url="+request.Referer())
				fmt.Fprintf(writer, "Attached image from %s deleted successfully
\n<meta http-equiv=\"refresh\" content=\"1;url="+config.DocumentRoot+"/"+board+"/\">", post)
			} else {

				// delete the post
				_, err = db.Exec("UPDATE `" + config.DBprefix + "posts` SET `deleted_timestamp` = '" + getSQLDateTime() + "' WHERE `id` = " + post)
				if parent_id == 0 {
					err = os.Remove(path.Join(config.DocumentRoot, board, "/res/"+post+".html"))
				} else {
					_board, _ := getBoardArr("`id` = " + strconv.Itoa(board_id))
					buildBoardPages(&_board[0])
				}

				// if the deleted post is actually a thread, delete its posts
				_, _ = db.Exec("UPDATE `" + config.DBprefix + "posts` SET `deleted_timestamp` = '" + getSQLDateTime() + "' WHERE `parentid` = " + post)
				if err != nil {
					server.ServeErrorPage(writer, err.Error())
					return
				}

				// delete the
				var deleted_filename string
				err = db.QueryRow("SELECT `filename` FROM `" + config.DBprefix + "posts` WHERE `id` = " + post + " AND `filename` != ''").Scan(&deleted_filename)
				if err == nil {
					os.Remove(path.Join(config.DocumentRoot, board, "/src/", deleted_filename))
					os.Remove(path.Join(config.DocumentRoot, board, "/thumb/", strings.Replace(deleted_filename, ".", "t.", -1)))
					os.Remove(path.Join(config.DocumentRoot, board, "/thumb/", strings.Replace(deleted_filename, ".", "c.", -1)))
				}

				err = db.QueryRow("SELECT `filename` FROM `" + config.DBprefix + "posts` WHERE `parentid` = " + post + " AND `filename` != ''").Scan(&deleted_filename)
				if err == nil {
					os.Remove(path.Join(config.DocumentRoot, board, "/src/", deleted_filename))
					os.Remove(path.Join(config.DocumentRoot, board, "/thumb/", strings.Replace(deleted_filename, ".", "t.", -1)))
					os.Remove(path.Join(config.DocumentRoot, board, "/thumb/", strings.Replace(deleted_filename, ".", "c.", -1)))
				}

				buildBoards(false, board_id)

				writer.Header().Add("refresh", "3;url="+request.Referer())
				fmt.Fprintf(writer, "%s deleted successfully\n
", post)
			}
		}
	}
}
<|endoftext|>


<commit_before>package hal

import (
	"bytes"
	"log"
	"strings"
)

// provides a persistent configuration store

// Order of precendence for prefs:
// user -> channel -> broker -> plugin -> global -> default

// PREFS_TABLE contains the SQL to create the prefs table
// key field is called pkey because key is a reserved word
const PREFS_TABLE = `
CREATE TABLE IF NOT EXISTS prefs (
	 user    VARCHAR(32) DEFAULT "",
	 channel VARCHAR(32) DEFAULT "",
	 broker  VARCHAR(32) DEFAULT "",
	 plugin  VARCHAR(32) DEFAULT "",
	 pkey    VARCHAR(32) NOT NULL,
	 value   TEXT,
	 PRIMARY KEY(user, channel, broker, plugin, pkey)
)`

/*
   -- test data, will remove once there are automated tests
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "", "", "", "foo", "user");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "CORE", "", "", "foo",
   "user-channel");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "CORE", "slack", "", "foo",
   "user-channel-broker");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "CORE", "slack", "uptime",
   "foo", "user-channel-broker-plugin");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "", "slack", "uptime", "foo",
   "user-broker-plugin");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "CORE", "", "uptime",
   "foo", "user-channel-plugin");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "", "", "uptime", "foo",
   "user-plugin");
*/

// !prefs list --scope plugin --plugin autoresponder
// !prefs get --scope channel --plugin autoresponder --channel CORE --key timezone
// !prefs set --scope user --plugin autoresponder --channel CORE

// Pref is a key/value pair associated with a combination of user, plugin,
// borker, or channel.
type Pref struct {
	User    string
	Plugin  string
	Broker  string
	Channel string
	Key     string
	Value   string
	Default string
	Success bool
	Error   error
}

type Prefs []*Pref

// GetPref will retreive the most-specific preference from pref
// storage using the parameters provided. This is a bit like pattern
// matching. If no match is found, the provided default is returned.
// TODO: explain this better
func GetPref(user, broker, channel, plugin, key, def string) Pref {
	pref := Pref{
		User:    user,
		Channel: channel,
		Broker:  broker,
		Plugin:  plugin,
		Key:     key,
		Default: def,
	}

	up := pref.Get()
	if up.Success {
		return up
	}

	// no match, return the default
	pref.Value = def
	return pref
}

// GetPrefs retrieves a set of preferences from the database. The
// settings are matched exactly on user,broker,channel,plugin.
// e.g. GetPrefs("", "", "", "uptime") would get only records that
// have user/broker/channel set to the empty string and channel
// set to "uptime". A record with user "pford" and plugin "uptime"
// would not be included.
func GetPrefs(user, broker, channel, plugin string) Prefs {
	pref := Pref{
		User:    user,
		Broker:  broker,
		Channel: channel,
		Plugin:  plugin,
	}
	return pref.get()
}

// FindPrefs gets all records that match any of the inputs that are
// not empty strings. (hint: user="x", broker="y"; WHERE user=? OR broker=?)
func FindPrefs(user, broker, channel, plugin, key string) Prefs {
	pref := Pref{
		User:    user,
		Broker:  broker,
		Channel: channel,
		Plugin:  plugin,
		Key:     key,
	}
	return pref.Find()
}

func GetUserPrefs(user string) Prefs {
	pref := Pref{}
	pref.User = user
	return pref.get()
}

func GetChannelPrefs(channel string) Prefs {
	pref := Pref{}
	pref.Channel = channel
	return pref.get()
}

func GetBrokerPrefs(broker string) Prefs {
	pref := Pref{}
	pref.Broker = broker
	return pref.get()
}

func GetPluginPrefs(plugin string) Prefs {
	pref := Pref{}
	pref.Plugin = plugin
	return pref.get()
}

// Get retrieves a value from the database. If the database returns
// an error, Success will be false and the Error field will be populated.
func (in *Pref) Get() Pref {
	prefs := in.get()

	if len(prefs) == 1 {
		return *prefs[0]
	} else if len(prefs) > 1 {
		panic("TOO MANY PREFS")
	} else if len(prefs) == 0 {
		out := *in
		out.Success = false
		return out
	}

	panic("BUG: should be impossible to reach this point")
}

func (in *Pref) get() Prefs {
	SqlInit(PREFS_TABLE)

	sql := `SELECT user,channel,broker,plugin,pkey,value
	        FROM prefs
	        WHERE user=?
			  AND channel=?
			  AND broker=?
			  AND plugin=?`
	params := []interface{}{&in.User, &in.Channel, &in.Broker, &in.Plugin}

	// only query by key if it's specified, otherwise get all keys for the selection
	if in.Key != "" {
		sql += " AND pkey=?"
		params = append(params, &in.Key)
	}

	db := SqlDB()

	rows, err := db.Query(sql, params...)
	if err != nil {
		log.Printf("Returning default due to SQL query failure: %s", err)
		return Prefs{}
	}

	defer rows.Close()

	out := make(Prefs, 0)

	for rows.Next() {
		p := *in

		err := rows.Scan(&p.User, &p.Channel, &p.Broker, &p.Plugin, &p.Key, &p.Value)

		if err != nil {
			log.Printf("Returning default due to row iteration failure: %s", err)
			p.Success = false
			p.Value = in.Default
			p.Error = err
		}

		out = append(out, &p)
	}

	return out
}

// Set writes the value and returns a new struct with the new value.
func (in *Pref) Set() Pref {
	db := SqlDB()
	SqlInit(PREFS_TABLE)

	sql := `INSERT INTO prefs
						(value,user,channel,broker,plugin,pkey)
			VALUES (?,?,?,?,?,?)
			ON DUPLICATE KEY
			UPDATE value=?, user=?, channel=?, broker=?, plugin=?, pkey=?`

	params := []interface{}{
		&in.Value, &in.User, &in.Channel, &in.Broker, &in.Plugin, &in.Key,
		&in.Value, &in.User, &in.Channel, &in.Broker, &in.Plugin, &in.Key,
	}

	_, err := db.Exec(sql, params...)
	if err != nil {
		out := *in
		out.Success = false
		out.Error = err
		return out
	}

	return in.Get()
}

// Find retrieves all preferences from the database that match any field in the
// handle's fields.
// Unlike Get(), empty string fields are not included in the (generated) query
// so it can potentially match a lot of rows.
// Returns an empty list and logs upon errors.
func (p Pref) Find() Prefs {
	SqlInit(PREFS_TABLE)

	fields := make([]string, 0)
	params := make([]interface{}, 0)

	if p.User != "" {
		fields = append(fields, "user=?")
		params = append(params, p.User)
	}

	if p.Channel != "" {
		fields = append(fields, "channel=?")
		params = append(params, p.Channel)
	}

	if p.Broker != "" {
		fields = append(fields, "broker=?")
		params = append(params, p.Broker)
	}

	if p.Plugin != "" {
		fields = append(fields, "plugin=?")
		params = append(params, p.Plugin)
	}

	if p.Key != "" {
		fields = append(fields, "pkey=?")
		params = append(params, p.Key)
	}

	q := bytes.NewBufferString("SELECT user,channel,broker,plugin,pkey,value\n")
	q.WriteString("FROM prefs\n")

	// TODO: maybe it's silly to make it easy for Find() to get all preferences
	// but let's cross that bridge when we come to it
	if len(fields) > 0 {
		q.WriteString("\nWHERE ")
		// might make sense to add a param to this func to make it easy to
		// switch this between AND/OR for unions/intersections
		q.WriteString(strings.Join(fields, "\n  OR "))
	}

	// TODO: add deterministic ordering at query time

	db := SqlDB()
	out := make(Prefs, 0)
	rows, err := db.Query(q.String(), params...)
	if err != nil {
		log.Println(q.String())
		log.Printf("Query failed: %s", err)
		return out
	}
	defer rows.Close()

	for rows.Next() {
		row := Pref{}
		err = rows.Scan(&row.User, &row.Channel, &row.Broker, &row.Plugin, &row.Key, &row.Value)
		// improbable in practice - follows previously mentioned conventions for errors
		if err != nil {
			log.Printf("Fetching a row failed: %s\n", err)
			row.Error = err
			row.Success = false
			row.Value = p.Default
		} else {
			row.Error = nil
			row.Success = true
		}

		out = append(out, &row)
	}

	return out
}

// User filters the preference list by user, returning a new Prefs
// e.g. uprefs = prefs.User("adent")
func (prefs Prefs) User(user string) Prefs {
	out := make(Prefs, 0)

	for _, pref := range prefs {
		if pref.User == user {
			out = append(out, pref)
		}
	}

	return out
}

// Channel filters the preference list by channel, returning a new Prefs
// e.g. instprefs = prefs.Channel("magrathea").Plugin("uptime").Broker("slack")
func (prefs Prefs) Channel(channel string) Prefs {
	out := make(Prefs, 0)

	for _, pref := range prefs {
		if pref.Channel == channel {
			out = append(out, pref)
		}
	}

	return out
}

// Broker filters the preference list by broker, returning a new Prefs
func (prefs Prefs) Broker(broker string) Prefs {
	out := make(Prefs, 0)

	for _, pref := range prefs {
		if pref.Broker == broker {
			out = append(out, pref)
		}
	}

	return out
}

// Plugin filters the preference list by plugin, returning a new Prefs
func (prefs Prefs) Plugin(plugin string) Prefs {
	out := make(Prefs, 0)

	for _, pref := range prefs {
		if pref.Plugin == plugin {
			out = append(out, pref)
		}
	}

	return out
}

// ready to hand off to e.g. hal.AsciiTable()
func (prefs Prefs) Table() [][]string {
	out := make([][]string, 1)
	out[0] = []string{"User", "Channel", "Broker", "Plugin", "Key", "Value"}

	for _, pref := range prefs {
		m := []string{
			pref.User,
			pref.Channel,
			pref.Broker,
			pref.Plugin,
			pref.Key,
			pref.Value,
		}

		out = append(out, m)
	}

	return out
}
<commit_msg>add GetPrefs, String, and Delete methods, fix errors<commit_after>package hal

import (
	"bytes"
	"fmt"
	"log"
	"strings"
)

// provides a persistent configuration store

// Order of precendence for prefs:
// user -> channel -> broker -> plugin -> global -> default

// PREFS_TABLE contains the SQL to create the prefs table
// key field is called pkey because key is a reserved word
const PREFS_TABLE = `
CREATE TABLE IF NOT EXISTS prefs (
	 user    VARCHAR(32) DEFAULT "",
	 channel VARCHAR(32) DEFAULT "",
	 broker  VARCHAR(32) DEFAULT "",
	 plugin  VARCHAR(32) DEFAULT "",
	 pkey    VARCHAR(32) NOT NULL,
	 value   TEXT,
	 PRIMARY KEY(user, channel, broker, plugin, pkey)
)`

/*
   -- test data, will remove once there are automated tests
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "", "", "", "foo", "user");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "CORE", "", "", "foo",
   "user-channel");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "CORE", "slack", "", "foo",
   "user-channel-broker");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "CORE", "slack", "uptime",
   "foo", "user-channel-broker-plugin");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "", "slack", "uptime", "foo",
   "user-broker-plugin");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "CORE", "", "uptime",
   "foo", "user-channel-plugin");
   INSERT INTO prefs (user,channel,broker,plugin,pkey,value) VALUES ("tobert", "", "", "uptime", "foo",
   "user-plugin");
*/

// !prefs list --scope plugin --plugin autoresponder
// !prefs get --scope channel --plugin autoresponder --channel CORE --key timezone
// !prefs set --scope user --plugin autoresponder --channel CORE

// Pref is a key/value pair associated with a combination of user, plugin,
// borker, or channel.
type Pref struct {
	User    string
	Plugin  string
	Broker  string
	Channel string
	Key     string
	Value   string
	Default string
	Success bool
	Error   error
}

type Prefs []*Pref

// GetPref will retreive the most-specific preference from pref
// storage using the parameters provided. This is a bit like pattern
// matching. If no match is found, the provided default is returned.
// TODO: explain this better
func GetPref(user, broker, channel, plugin, key, def string) Pref {
	pref := Pref{
		User:    user,
		Channel: channel,
		Broker:  broker,
		Plugin:  plugin,
		Key:     key,
		Default: def,
	}

	up := pref.Get()
	if up.Success {
		return up
	}

	// no match, return the default
	pref.Value = def
	return pref
}

// GetPrefs retrieves a set of preferences from the database. The
// settings are matched exactly on user,broker,channel,plugin.
// e.g. GetPrefs("", "", "", "uptime") would get only records that
// have user/broker/channel set to the empty string and channel
// set to "uptime". A record with user "pford" and plugin "uptime"
// would not be included.
func GetPrefs(user, broker, channel, plugin string) Prefs {
	pref := Pref{
		User:    user,
		Broker:  broker,
		Channel: channel,
		Plugin:  plugin,
	}
	return pref.get()
}

// FindPrefs gets all records that match any of the inputs that are
// not empty strings. (hint: user="x", broker="y"; WHERE user=? OR broker=?)
func FindPrefs(user, broker, channel, plugin, key string) Prefs {
	pref := Pref{
		User:    user,
		Broker:  broker,
		Channel: channel,
		Plugin:  plugin,
		Key:     key,
	}
	return pref.Find()
}

func GetUserPrefs(user string) Prefs {
	pref := Pref{}
	pref.User = user
	return pref.get()
}

func GetChannelPrefs(channel string) Prefs {
	pref := Pref{}
	pref.Channel = channel
	return pref.get()
}

func GetBrokerPrefs(broker string) Prefs {
	pref := Pref{}
	pref.Broker = broker
	return pref.get()
}

func GetPluginPrefs(plugin string) Prefs {
	pref := Pref{}
	pref.Plugin = plugin
	return pref.get()
}

// Get retrieves a value from the database. If the database returns
// an error, Success will be false and the Error field will be populated.
func (in *Pref) Get() Pref {
	prefs := in.get()

	if len(prefs) == 1 {
		return *prefs[0]
	} else if len(prefs) > 1 {
		panic("TOO MANY PREFS")
	} else if len(prefs) == 0 {
		out := *in
		// only set success to false if there is also an error
		// queries with 0 rows are successful
		if out.Error != nil {
			out.Success = false
		} else {
			out.Success = true
			out.Value = out.Default
		}
		return out
	}

	panic("BUG: should be impossible to reach this point")
}

// GetPrefs returns all preferences that match the fields set in the handle.
func (in *Pref) GetPrefs() Prefs {
	return in.get()
}

func (in *Pref) get() Prefs {
	SqlInit(PREFS_TABLE)

	sql := `SELECT user,channel,broker,plugin,pkey,value
	        FROM prefs
	        WHERE user=?
			  AND channel=?
			  AND broker=?
			  AND plugin=?`
	params := []interface{}{&in.User, &in.Channel, &in.Broker, &in.Plugin}

	// only query by key if it's specified, otherwise get all keys for the selection
	if in.Key != "" {
		sql += " AND pkey=?"
		params = append(params, &in.Key)
	}

	db := SqlDB()

	rows, err := db.Query(sql, params...)
	if err != nil {
		log.Printf("Returning default due to SQL query failure: %s", err)
		return Prefs{}
	}

	defer rows.Close()

	out := make(Prefs, 0)

	for rows.Next() {
		p := *in

		err := rows.Scan(&p.User, &p.Channel, &p.Broker, &p.Plugin, &p.Key, &p.Value)

		if err != nil {
			log.Printf("Returning default due to row iteration failure: %s", err)
			p.Success = false
			p.Value = in.Default
			p.Error = err
		} else {
			p.Success = true
			p.Error = nil
		}

		out = append(out, &p)
	}

	return out
}

// Set writes the value and returns a new struct with the new value.
func (in *Pref) Set() error {
	db := SqlDB()
	SqlInit(PREFS_TABLE)

	sql := `INSERT INTO prefs
						(value,user,channel,broker,plugin,pkey)
			VALUES (?,?,?,?,?,?)
			ON DUPLICATE KEY
			UPDATE value=?, user=?, channel=?, broker=?, plugin=?, pkey=?`

	params := []interface{}{
		&in.Value, &in.User, &in.Channel, &in.Broker, &in.Plugin, &in.Key,
		&in.Value, &in.User, &in.Channel, &in.Broker, &in.Plugin, &in.Key,
	}

	_, err := db.Exec(sql, params...)
	if err != nil {
		log.Printf("Pref.Set() write failed: %s", err)
		return err
	}

	return nil
}

// Set writes the value and returns a new struct with the new value.
func (in *Pref) Delete() error {
	db := SqlDB()
	SqlInit(PREFS_TABLE)

	sql := `DELETE FROM prefs
			WHERE user=?
			  AND channel=?
			  AND broker=?
			  AND plugin=?
			  AND pkey=?`

	// TODO: verify only one row was deleted
	_, err := db.Exec(sql, &in.User, &in.Channel, &in.Broker, &in.Plugin, &in.Key)
	if err != nil {
		log.Printf("Pref.Delete() write failed: %s", err)
		return err
	}

	return nil
}

// Find retrieves all preferences from the database that match any field in the
// handle's fields.
// Unlike Get(), empty string fields are not included in the (generated) query
// so it can potentially match a lot of rows.
// Returns an empty list and logs upon errors.
func (p Pref) Find() Prefs {
	SqlInit(PREFS_TABLE)

	fields := make([]string, 0)
	params := make([]interface{}, 0)

	if p.User != "" {
		fields = append(fields, "user=?")
		params = append(params, p.User)
	}

	if p.Channel != "" {
		fields = append(fields, "channel=?")
		params = append(params, p.Channel)
	}

	if p.Broker != "" {
		fields = append(fields, "broker=?")
		params = append(params, p.Broker)
	}

	if p.Plugin != "" {
		fields = append(fields, "plugin=?")
		params = append(params, p.Plugin)
	}

	if p.Key != "" {
		fields = append(fields, "pkey=?")
		params = append(params, p.Key)
	}

	q := bytes.NewBufferString("SELECT user,channel,broker,plugin,pkey,value\n")
	q.WriteString("FROM prefs\n")

	// TODO: maybe it's silly to make it easy for Find() to get all preferences
	// but let's cross that bridge when we come to it
	if len(fields) > 0 {
		q.WriteString("\nWHERE ")
		// might make sense to add a param to this func to make it easy to
		// switch this between AND/OR for unions/intersections
		q.WriteString(strings.Join(fields, "\n  OR "))
	}

	// TODO: add deterministic ordering at query time

	db := SqlDB()
	out := make(Prefs, 0)
	rows, err := db.Query(q.String(), params...)
	if err != nil {
		log.Println(q.String())
		log.Printf("Query failed: %s", err)
		return out
	}
	defer rows.Close()

	for rows.Next() {
		row := Pref{}
		err = rows.Scan(&row.User, &row.Channel, &row.Broker, &row.Plugin, &row.Key, &row.Value)
		// improbable in practice - follows previously mentioned conventions for errors
		if err != nil {
			log.Printf("Fetching a row failed: %s\n", err)
			row.Error = err
			row.Success = false
			row.Value = p.Default
		} else {
			row.Error = nil
			row.Success = true
		}

		out = append(out, &row)
	}

	return out
}

// User filters the preference list by user, returning a new Prefs
// e.g. uprefs = prefs.User("adent")
func (prefs Prefs) User(user string) Prefs {
	out := make(Prefs, 0)

	for _, pref := range prefs {
		if pref.User == user {
			out = append(out, pref)
		}
	}

	return out
}

// Channel filters the preference list by channel, returning a new Prefs
// e.g. instprefs = prefs.Channel("magrathea").Plugin("uptime").Broker("slack")
func (prefs Prefs) Channel(channel string) Prefs {
	out := make(Prefs, 0)

	for _, pref := range prefs {
		if pref.Channel == channel {
			out = append(out, pref)
		}
	}

	return out
}

// Broker filters the preference list by broker, returning a new Prefs
func (prefs Prefs) Broker(broker string) Prefs {
	out := make(Prefs, 0)

	for _, pref := range prefs {
		if pref.Broker == broker {
			out = append(out, pref)
		}
	}

	return out
}

// Plugin filters the preference list by plugin, returning a new Prefs
func (prefs Prefs) Plugin(plugin string) Prefs {
	out := make(Prefs, 0)

	for _, pref := range prefs {
		if pref.Plugin == plugin {
			out = append(out, pref)
		}
	}

	return out
}

// ready to hand off to e.g. hal.AsciiTable()
func (prefs Prefs) Table() [][]string {
	out := make([][]string, 1)
	out[0] = []string{"User", "Channel", "Broker", "Plugin", "Key", "Value"}

	for _, pref := range prefs {
		m := []string{
			pref.User,
			pref.Channel,
			pref.Broker,
			pref.Plugin,
			pref.Key,
			pref.Value,
		}

		out = append(out, m)
	}

	return out
}

func (p *Pref) String() string {
	return fmt.Sprintf(`Pref{
	User:    %q,
	Channel: %q,
	Broker:  %q,
	Plugin:  %q,
	Key:     %q,
	Value:   %q,
	Default: %q,
	Success: %t,
	Error:   %v,
}`, p.User, p.Channel, p.Broker, p.Plugin, p.Key, p.Value, p.Default, p.Success, p.Error)

}

func (p *Prefs) String() string {
	data := p.Table()
	return AsciiTable(data[0], data[1:])
}
<|endoftext|>


<commit_before>package hapi

import (
    "testing"
    "net/http"
//    "net/http/httptest"
)

type mockResponseWriter struct{}

func (m *mockResponseWriter) Header() (h http.Header) {
    return http.Header{}
}

func (m *mockResponseWriter) Write(p []byte) (n int, err error) {
    return len(p), nil
}

func (m *mockResponseWriter) WriteString(s string) (n int, err error) {
    return len(s), nil
}

func (m *mockResponseWriter) WriteHeader(int) {}

func TestNew(t *testing.T) {
    h := New()
    if h.Router == nil {
        t.Fatal("h.Router not initilized")
    }
}

type typeTest struct{
    Accept      string
    Handlers    map[string]Handle
    Result      string
    HandlerID   string
}

func NewTest(accept,result,id string) *typeTest {
    return &typeTest{
        accept,
        make(map[string]Handle),
        result,
        id,
    }
}

func (testData *typeTest) AddHandler(ctype,id string) {
    testData.Handlers[ctype] = func(c *Context) { c.Stash["id"] = id }
}

func (testData *typeTest) DoTests(t *testing.T) {
    negType,negHandler := TypeNegotiator(testData.Accept,testData.Handlers)
    if negType != testData.Result {
        t.Fatalf("TypeNegotiator returned '%s', expected '%s'\n", negType, testData.Result)
    }
    context := &Context{}
    context.Stash = make(map[string]interface{})
    if negHandler != nil {
        negHandler( context )
    }
    var id string
    if context.Stash["id"] != nil {
        id = context.Stash["id"].(string)
    }
    if id != testData.HandlerID {
        t.Fatalf("Handler identified itself as '%s', expected '%s'\n", id, testData.HandlerID)
    }
}

// Tests 'text/html' handler
func TestTypeNegotiator(t *testing.T) {
    // Accept: matches exactly
    testData := NewTest("text/html","text/html","1")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: Doesn't match
    testData = NewTest("text/plain","","")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: doesn't exist
    testData = NewTest("","text/html","1")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: */*
    testData = NewTest("*/*","text/html","1")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: text/*
    testData = NewTest("text/*","text/html","1")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: foo/*
    testData = NewTest("text/plain","","")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)
}

// Tests for '*/*' handler
func TestTypeNegotiator2(t *testing.T) {
    // Accept: text/html
    testData := NewTest("text/html","text/html","1")
    testData.AddHandler("*/*","1")
    testData.DoTests(t)

    testData = NewTest("*/*","*/*","1")
    testData.AddHandler("*/*","1")
    testData.DoTests(t)

    testData = NewTest("","*/*","1")
    testData.AddHandler("*/*","1")
    testData.DoTests(t)
}

// Tests for two routers 1: "text/html", 2: "*/*"
func TestTypeNegotiator3(t *testing.T) {
    testData := NewTest("text/html","text/html","1")
    testData.AddHandler("text/html","1")
    testData.AddHandler("*/*","2")
    testData.DoTests(t)

    testData = NewTest("text/html","text/html","2")
    testData.AddHandler("text/plain","1")
    testData.AddHandler("*/*","2")
    testData.DoTests(t)

    testData = NewTest("text/*","text/html","1")
    testData.AddHandler("text/html","1")
    testData.AddHandler("*/*","2")
    testData.DoTests(t)
}

// Two routers: 1: text/html, 2: text/plain
func TestTypeNegotiator4(t *testing.T) {
    testData := NewTest("text/html","text/html","1")
    testData.AddHandler("text/html","1")
    testData.AddHandler("text/plain","2")
    testData.DoTests(t)

    testData = NewTest("text/plain","text/plain","1")
    testData.AddHandler("text/plain","1")
    testData.AddHandler("text/html","2")
    testData.DoTests(t)

    testData = NewTest("text/*","text/html","1")
    testData.AddHandler("text/*","1")
    testData.AddHandler("text/plain","2")
    testData.DoTests(t)

}
<commit_msg>Add test names<commit_after>package hapi

import (
    "testing"
    "net/http"
//    "net/http/httptest"
)

type mockResponseWriter struct{}

func (m *mockResponseWriter) Header() (h http.Header) {
    return http.Header{}
}

func (m *mockResponseWriter) Write(p []byte) (n int, err error) {
    return len(p), nil
}

func (m *mockResponseWriter) WriteString(s string) (n int, err error) {
    return len(s), nil
}

func (m *mockResponseWriter) WriteHeader(int) {}

func TestNew(t *testing.T) {
    h := New()
    if h.Router == nil {
        t.Fatal("h.Router not initilized")
    }
}

type typeTest struct{
    Name        string
    Accept      string
    Handlers    map[string]Handle
    Result      string
    HandlerID   string
}

func NewTest(name,accept,result,id string) *typeTest {
    return &typeTest{
        name,
        accept,
        make(map[string]Handle),
        result,
        id,
    }
}

func (testData *typeTest) AddHandler(ctype,id string) {
    testData.Handlers[ctype] = func(c *Context) { c.Stash["id"] = id }
}

func (testData *typeTest) DoTests(t *testing.T) {
    negType,negHandler := TypeNegotiator(testData.Accept,testData.Handlers)
    if negType != testData.Result {
        t.Fatalf("%s: TypeNegotiator returned '%s', expected '%s'\n", testData.Name, negType, testData.Result)
    }
    context := &Context{}
    context.Stash = make(map[string]interface{})
    if negHandler != nil {
        negHandler( context )
    }
    var id string
    if context.Stash["id"] != nil {
        id = context.Stash["id"].(string)
    }
    if id != testData.HandlerID {
        t.Fatalf("%s: Handler identified itself as '%s', expected '%s'\n", testData.Name, id, testData.HandlerID)
    }
}

// Tests 'text/html' handler
func TestTypeNegotiator(t *testing.T) {
    // Accept: matches exactly
    testData := NewTest("test/html 1","text/html","text/html","1")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: Doesn't match
    testData = NewTest("test/html 2","text/plain","","")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: doesn't exist
    testData = NewTest("test/html 3","","text/html","1")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: */*
    testData = NewTest("test/html 4","*/*","text/html","1")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: text/*
    testData = NewTest("test/html 5","text/*","text/html","1")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)

    // Accept: foo/*
    testData = NewTest("test/html 6","text/plain","","")
    testData.AddHandler("text/html","1")
    testData.DoTests(t)
}

// Tests for '*/*' handler
func TestTypeNegotiator2(t *testing.T) {
    // Accept: text/html
    testData := NewTest("*/* 1","text/html","text/html","1")
    testData.AddHandler("*/*","1")
    testData.DoTests(t)

    testData = NewTest("*/* 2","*/*","*/*","1")
    testData.AddHandler("*/*","1")
    testData.DoTests(t)

    testData = NewTest("*/* 3","","*/*","1")
    testData.AddHandler("*/*","1")
    testData.DoTests(t)
}

// Tests for two routers 1: "text/html", 2: "*/*"
func TestTypeNegotiator3(t *testing.T) {
    testData := NewTest("text/html + */* 1","text/html","text/html","1")
    testData.AddHandler("text/html","1")
    testData.AddHandler("*/*","2")
    testData.DoTests(t)

    testData = NewTest("text/html + */* 2","text/html","text/html","2")
    testData.AddHandler("text/plain","1")
    testData.AddHandler("*/*","2")
    testData.DoTests(t)

    testData = NewTest("text/html + */* 3","text/*","text/*","2")
    testData.AddHandler("text/html","1")
    testData.AddHandler("*/*","2")
    testData.DoTests(t)
}

// Two routers: 1: text/html, 2: text/plain
func TestTypeNegotiator4(t *testing.T) {
    testData := NewTest("text/html + text/plain 1","text/html","text/html","1")
    testData.AddHandler("text/html","1")
    testData.AddHandler("text/plain","2")
    testData.DoTests(t)

    testData = NewTest("text/html + text/plain 2","text/plain","text/plain","2")
    testData.AddHandler("text/html","1")
    testData.AddHandler("text/plain","2")
    testData.DoTests(t)

    testData = NewTest("text/html + text/plain 3","text/*","text/*","1")
    testData.AddHandler("text/*","1")
    testData.AddHandler("text/plain","2")
    testData.DoTests(t)
}

// Three routers: 1: */*, 2: text/html, 3: text/plain
func TestTypeNegotiator5(t *testing.T) {
    testData := NewTest("*/* + text/html + text/plain 1", "text/html","text/html","2")
    testData.AddHandler("*/*","1")
    testData.AddHandler("text/html","2")
    testData.AddHandler("text/plain","3")

    testData = NewTest("*/* + text/html + text/plain 2", "text/html","text/foo","1")
    testData.AddHandler("*/*","1")
    testData.AddHandler("text/html","2")
    testData.AddHandler("text/plain","3")

    testData = NewTest("*/* + text/html + text/plain 3", "text/*","text/*","2")
    testData.AddHandler("*/*","1")
    testData.AddHandler("text/html","2")
    testData.AddHandler("text/plain","3")

}
<|endoftext|>


<commit_before>package config

import "testing"

func TestTransformSuccess(t *testing.T) {
	config := createSmallTestData()
	if config.RunName != "name" {
		t.Error("Can't transform JSON to struct")
	}
	if config.Test.Filter != "" {
		t.Error("JSON transformation to struct is incorrect")
	}
}

func TestSizeOfArrays(t *testing.T) {
	config := createSmallTestData()
	if len(config.AdditionalData.AuxiliaryApps) != 0 {
		t.Error("Size should be 0")
	}
	if len(config.Test.Parameters) != 0 {
		t.Error("Size should be 0")
	}
}

func TestContainMapData(t *testing.T) {
	config := createMapTestData()
	if len(config.Test.Parameters) != 1 {
		t.Error("Size of map should be 1")
	}
	if config.Test.Parameters["key"] != "value" {
		t.Error("Result should be equal to 'value'")
	}
}

func TestDontContainMapData(t *testing.T) {
	config := createSmallTestData()
	if len(config.Test.Parameters) != 0 {
		t.Error("Size of map should be 0")
	}
}

func TestGetUploadUnexisted(t *testing.T) {
	res := GetUploadTypeForTest("UIAUTOMATION")
	if res != "" {
		t.Error("For unexisted type of input result should be blank string")
	}
}

func TestGetUploadPositive(t *testing.T) {
	res := GetUploadTypeForTest("APPIUM_WEB_JAVA_JUNIT")
	if res != "APPIUM_WEB_JAVA_JUNIT_TEST_PACKAGE" {
		t.Error("For unexisted type of input result should be blank string")
	}
}

func createSmallTestData() RunConfig {
	json := []byte(`{"runName":"name"}`)
	return Transform(json)
}

func createMapTestData() RunConfig {
	json := []byte(`{"runName":"name","test":{"parameters":{"key":"value"}}}`)
	return Transform(json)
}
<commit_msg>Fix errors messages on test failing<commit_after>package config

import "testing"

func TestTransformSuccess(t *testing.T) {
	config := createSmallTestData()
	if config.RunName != "name" {
		t.Error("Can't transform JSON to struct")
	}
	if config.Test.Filter != "" {
		t.Error("JSON transformation to struct is incorrect")
	}
}

func TestSizeOfArrays(t *testing.T) {
	config := createSmallTestData()
	if len(config.AdditionalData.AuxiliaryApps) != 0 {
		t.Error("Size of AdditionalData.AuxiliaryApps should be 0")
	}
	if len(config.Test.Parameters) != 0 {
		t.Error("Size of Test.Parameters should be 0")
	}
}

func TestContainMapData(t *testing.T) {
	config := createMapTestData()
	if len(config.Test.Parameters) != 1 {
		t.Error("Size of Test.Parameters should be 1")
	}
	if config.Test.Parameters["key"] != "value" {
		t.Error("Test.Parameters['key'] should return 'value'")
	}
}

func TestDontContainMapData(t *testing.T) {
	config := createSmallTestData()
	if len(config.Test.Parameters) != 0 {
		t.Error("Size of Test.Parameters should be 0")
	}
}

func TestGetUploadUnexisted(t *testing.T) {
	res := GetUploadTypeForTest("UIAUTOMATION")
	if res != "" {
		t.Error("For unexisted type of input result should be blank string")
	}
}

func TestGetUploadPositive(t *testing.T) {
	res := GetUploadTypeForTest("APPIUM_WEB_JAVA_JUNIT")
	if res != "APPIUM_WEB_JAVA_JUNIT_TEST_PACKAGE" {
		t.Error("For unexisted type of input result should be blank string")
	}
}

func createSmallTestData() RunConfig {
	json := []byte(`{"runName":"name"}`)
	return Transform(json)
}

func createMapTestData() RunConfig {
	json := []byte(`{"runName":"name","test":{"parameters":{"key":"value"}}}`)
	return Transform(json)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

var heartFreq = flag.Duration("heartbeat", 10*time.Second,
	"Heartbeat frequency")
var reconcileFreq = flag.Duration("reconcile", 24*time.Hour,
	"Reconciliation frequency")
var staleNodeFreq = flag.Duration("staleNodeCheck", 5*time.Minute,
	"How frequently to check for stale nodes.")
var staleNodeLimit = flag.Duration("staleNodeLimit", 15*time.Minute,
	"How long until we clean up nodes for being too stale")
var nodeCleanCount = flag.Int("nodeCleanCount", 1000,
	"How many blobs to clean up from a dead node per period")
var verifyWorkers = flag.Int("verifyWorkers", 4,
	"Number of object verification workers.")
var garbageCollectFreq = flag.Duration("gcFreq", 5*time.Minute,
	"How frequently to check dangling blobs.")
var maxStartupObjects = flag.Int("maxStartObjs", 1000,
	"Maximum number of objects to pull on start")
var maxStartupRepls = flag.Int("maxStartRepls", 3,
	"Blob replication limit for startup objects.")

type PeriodicJob struct {
	period time.Duration
	f      func() error
}

var periodicJobs = map[string]*PeriodicJob{
	"checkStaleNodes": &PeriodicJob{
		time.Minute * 5,
		checkStaleNodes,
	},
	"garbageCollectBlobs": &PeriodicJob{
		time.Minute * 5,
		garbageCollectBlobs,
	},
}

func adjustPeriodicJobs() error {
	periodicJobs["checkStaleNodes"].period = *staleNodeFreq
	periodicJobs["garbageCollectBlobs"].period = *garbageCollectFreq
	return nil
}

type JobMarker struct {
	Node    string    `json:"node"`
	Started time.Time `json:"started"`
	Type    string    `json:"type"`
}

// Run a named task if we know one hasn't in the last t seconds.
func runNamedGlobalTask(name string, t time.Duration, f func() error) bool {
	key := "/@" + name

	jm := JobMarker{
		Node:    serverId,
		Started: time.Now(),
		Type:    "job",
	}

	err := couchbase.Do(key, func(mc *memcached.Client, vb uint16) error {
		resp, err := mc.Add(vb, key, 0, int(t.Seconds()),
			mustEncode(&jm))
		if err != nil {
			return err
		}
		if resp.Status != gomemcached.SUCCESS {
			return fmt.Errorf("Wanted success, got %v", resp.Status)
		}
		return nil
	})

	if err == nil {
		err = f()
		if err != nil {
			log.Printf("Error running periodic task %#v: %v", name, err)
		}
		return true
	}

	return false
}

func heartbeat() {
	for {
		u, err := url.Parse(*couchbaseServer)
		c, err := net.Dial("tcp", u.Host)
		localAddr := ""
		if err == nil {
			localAddr = strings.Split(c.LocalAddr().String(), ":")[0]
			c.Close()
		}

		aboutMe := StorageNode{
			Addr:     localAddr,
			Type:     "node",
			Time:     time.Now().UTC(),
			BindAddr: *bindAddr,
			Hash:     *hashType,
		}

		err = couchbase.Set("/"+serverId, aboutMe)
		if err != nil {
			log.Printf("Failed to record a heartbeat: %v", err)
		}
		time.Sleep(*heartFreq)
	}
}

func verifyObjectHash(h string) error {
	fn := hashFilename(*root, h)
	f, err := os.Open(fn)
	if err != nil {
		return err
	}
	defer f.Close()

	sh := getHash()
	_, err = io.Copy(sh, f)
	if err != nil {
		return err
	}

	hstring := hex.EncodeToString(sh.Sum([]byte{}))
	if h != hstring {
		err = os.Remove(fn)
		if err != nil {
			log.Printf("Error removing corrupt file %v: %v", err)
		}
		return fmt.Errorf("Hash from disk of %v was %v", h, hstring)
	}
	return nil
}

func verifyWorker(ch chan os.FileInfo) {
	for info := range ch {
		err := verifyObjectHash(info.Name())
		if err == nil {
			recordBlobOwnership(info.Name(), info.Size())
		} else {
			log.Printf("Invalid hash for object %v found at verification: %v",
				info.Name(), err)
			removeBlobOwnershipRecord(info.Name(), serverId)
		}
	}
}

func reconcile() error {
	explen := getHash().Size() * 2

	vch := make(chan os.FileInfo)
	defer close(vch)

	for i := 0; i < *verifyWorkers; i++ {
		go verifyWorker(vch)
	}

	return filepath.Walk(*root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && !strings.HasPrefix(info.Name(), "tmp") &&
			len(info.Name()) == explen {

			vch <- info

			return err
		}
		return nil
	})
}

func reconcileLoop() {
	if *reconcileFreq == 0 {
		return
	}
	for {
		err := reconcile()
		if err != nil {
			log.Printf("Error in reconciliation loop: %v", err)
		}
		grabSomeData()
		time.Sleep(*reconcileFreq)
	}
}

func removeBlobOwnershipRecord(h, node string) {
	log.Printf("Cleaning up %v from %v", h, node)

	k := "/" + h
	err := couchbase.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.CAS(vb, k, func(in []byte) ([]byte, memcached.CasOp) {
			ownership := BlobOwnership{}
			err := json.Unmarshal(in, &ownership)
			if err == nil {
				if _, ok := ownership.Nodes[node]; !ok {
					// Skip it fast if we don't have it.
					return nil, memcached.CASQuit
				}
				delete(ownership.Nodes, node)
			} else {
				log.Printf("Error unmarhaling blob removal from %s: %v",
					in, err)
				return nil, memcached.CASQuit
			}

			var rv []byte
			op := memcached.CASStore

			if len(ownership.Nodes) == 0 {
				op = memcached.CASDelete
			} else {
				rv = mustEncode(&ownership)
			}

			return rv, op
		}, 0)
		return err
	})
	if err != nil && err != memcached.CASQuit {
		log.Printf("Error cleaning %v from %v: %v", node, h, err)
	}
}

func cleanupNode(node string) {
	log.Printf("Cleaning up node %v", node)
	vres, err := couchbase.View("cbfs", "node_blobs",
		map[string]interface{}{
			"key":    `"` + node + `"`,
			"limit":  *nodeCleanCount,
			"reduce": false,
			"stale":  false,
		})
	if err != nil {
		log.Printf("Error executing node_blobs view: %v", err)
		return
	}
	foundRows := 0
	for _, r := range vres.Rows {
		removeBlobOwnershipRecord(r.ID[1:], node)
		foundRows++
	}
	if foundRows == 0 && len(vres.Errors) == 0 {
		log.Printf("Removing node record: %v", node)
		err = couchbase.Delete("/" + node)
		if err != nil {
			log.Printf("Error deleting %v node record: %v", node, err)
		}
		err = couchbase.Delete("/" + node + "/r")
		if err != nil {
			log.Printf("Error deleting %v node counter: %v", node, err)
		}
	}
}

func checkStaleNodes() error {
	log.Printf("Checking stale nodes")
	vres, err := couchbase.View("cbfs", "nodes", map[string]interface{}{
		"stale": false})
	if err != nil {
		return err
	}
	for _, r := range vres.Rows {
		ks, ok := r.Key.(string)
		if !ok {
			log.Printf("Wrong key type returned from view: %#v", r)
			continue
		}
		t, err := time.Parse(time.RFC3339, ks)
		if err != nil {
			log.Printf("Error parsing time from %v", r)
			continue
		}
		d := time.Since(t)
		node := r.ID[1:]

		if d > *staleNodeLimit {
			if node == serverId {
				log.Printf("Would've cleaned up myself after %v",
					d)
				continue
			}
			log.Printf("  Node %v missed heartbeat schedule: %v", node, d)
			go cleanupNode(node)
		} else {
			log.Printf("%v is ok at %v", node, d)
		}
	}
	return nil
}

func garbageCollectBlobs() error {
	log.Printf("Garbage collecting blobs without any file references")

	viewRes := struct {
		Rows []struct {
			Key []string
		}
	}{}

	// we hit this view descending because we want file sorted before blob
	// the fact that we walk the list backwards hopefully not too awkward
	err := couchbase.ViewCustom("cbfs", "file_blobs",
		map[string]interface{}{
			"stale": false, "descending": true}, &viewRes)

	if err != nil {
		return err
	}

	lastBlob := ""
	count := 0
	for _, r := range viewRes.Rows {
		blobId := r.Key[0]
		typeFlag := r.Key[1]
		blobNode := r.Key[2]

		switch typeFlag {
		case "file":
			lastBlob = blobId
		case "blob":
			if blobId != lastBlob {
				go garbageCollectBlobFromNode(blobId, blobNode)
				count++
			}
		}

	}
	log.Printf("Scheduled %d blobs for deletion", count)
	return nil
}

func garbageCollectBlobFromNode(oid, sid string) {
	if sid == serverId {
		//local delete
		err := os.Remove(hashFilename(*root, oid))
		if err != nil {
			log.Printf("Error removing blob, already deleted? %v", err)
		}
	} else {
		//remote
		remote := StorageNode{}
		remotekey := "/" + sid
		err := couchbase.Get(remotekey, &remote)
		if err != nil {
			// will state node cleanup these records or should i?
			log.Printf("No record of this node")
			return
		}

		req, err := http.NewRequest("DELETE", remote.BlobURL(oid), nil)
		resp, err := http.DefaultClient.Do(req)

		if err != nil {
			log.Printf("Error deleting oid %s from node %s, %v",
				oid, sid, err)
			return
		}

		// non-obvious to me at first, but also with 404 we
		// should also remove the blob ownership
		if resp.StatusCode != 204 && resp.StatusCode != 404 {
			log.Printf("Unexpected status %v deleting %s from node %s",
				resp.Status, oid, sid)
			return
		}
	}
	removeBlobOwnershipRecord(oid, sid)
	log.Printf("Removed blob: %v from node %v", oid, sid)
}

type fetchSpec struct {
	oid  string
	node string
}

func dataInitFetchOne(h, u string) error {
	f, err := NewHashRecord(*root, h)
	if err != nil {
		return err
	}
	defer f.Close()

	resp, err := http.Get(u)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	h, l, err := f.Process(resp.Body)
	if err != nil {
		return err
	}
	return recordBlobOwnership(h, l)
}

func dataInitFetcher(nm map[string]StorageNode, ch <-chan fetchSpec) {
	for fs := range ch {
		node, found := nm[fs.node]
		if !found {
			log.Printf("couldn't find %v", fs.node)
			continue
		}
		log.Printf("Fetching %v from %v", fs.oid, node.BlobURL(fs.oid))
		err := dataInitFetchOne(fs.oid, node.BlobURL(fs.oid))
		if err != nil {
			log.Printf("Error fetching %v: %v", fs.oid, err)
		}
	}
}

func grabSomeData() {
	viewRes := struct {
		Rows []struct {
			Id  string
			Doc struct {
				Json struct {
					Nodes map[string]string
				}
			}
		}
	}{}

	// Find some less replicated docs to suck in.
	err := couchbase.ViewCustom("cbfs", "repcounts",
		map[string]interface{}{
			"reduce":       false,
			"include_docs": true,
			"limit":        *maxStartupObjects,
			"startkey":     1,
			"endkey":       1 + *maxStartupRepls,
		},
		&viewRes)

	if err != nil {
		log.Printf("Error finding docs to suck: %v", err)
		return
	}

	nl, err := findRemoteNodes()
	if err != nil {
		log.Printf("Error finding nodes: %v", err)
		return
	}
	nm := map[string]StorageNode{}

	for _, n := range nl {
		nm[n.name] = n
	}

	ch := make(chan fetchSpec, 1000)
	defer close(ch)

	for i := 0; i < 4; i++ {
		go dataInitFetcher(nm, ch)
	}

	for _, r := range viewRes.Rows {
		if _, ok := r.Doc.Json.Nodes[serverId]; !ok {
			for n := range r.Doc.Json.Nodes {
				if n != serverId {
					ch <- fetchSpec{r.Id[1:], n}
				}
			}
		}
	}
}

func runPeriodicJob(name string, job *PeriodicJob) {
	time.Sleep(time.Second * time.Duration(5+rand.Intn(60)))
	for {
		if runNamedGlobalTask(name, job.period, job.f) {
			log.Printf("Attempted job %v", name)
		} else {
			log.Printf("Didn't run job %v", name)
		}
		time.Sleep(job.period + time.Second)
	}
}

func runPeriodicJobs() {
	for n, j := range periodicJobs {
		go runPeriodicJob(n, j)
	}
}
<commit_msg>Use consolidated node list for checkStaleNodes<commit_after>package main

import (
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

var heartFreq = flag.Duration("heartbeat", 10*time.Second,
	"Heartbeat frequency")
var reconcileFreq = flag.Duration("reconcile", 24*time.Hour,
	"Reconciliation frequency")
var staleNodeFreq = flag.Duration("staleNodeCheck", 5*time.Minute,
	"How frequently to check for stale nodes.")
var staleNodeLimit = flag.Duration("staleNodeLimit", 15*time.Minute,
	"How long until we clean up nodes for being too stale")
var nodeCleanCount = flag.Int("nodeCleanCount", 1000,
	"How many blobs to clean up from a dead node per period")
var verifyWorkers = flag.Int("verifyWorkers", 4,
	"Number of object verification workers.")
var garbageCollectFreq = flag.Duration("gcFreq", 5*time.Minute,
	"How frequently to check dangling blobs.")
var maxStartupObjects = flag.Int("maxStartObjs", 1000,
	"Maximum number of objects to pull on start")
var maxStartupRepls = flag.Int("maxStartRepls", 3,
	"Blob replication limit for startup objects.")

type PeriodicJob struct {
	period time.Duration
	f      func() error
}

var periodicJobs = map[string]*PeriodicJob{
	"checkStaleNodes": &PeriodicJob{
		time.Minute * 5,
		checkStaleNodes,
	},
	"garbageCollectBlobs": &PeriodicJob{
		time.Minute * 5,
		garbageCollectBlobs,
	},
}

func adjustPeriodicJobs() error {
	periodicJobs["checkStaleNodes"].period = *staleNodeFreq
	periodicJobs["garbageCollectBlobs"].period = *garbageCollectFreq
	return nil
}

type JobMarker struct {
	Node    string    `json:"node"`
	Started time.Time `json:"started"`
	Type    string    `json:"type"`
}

// Run a named task if we know one hasn't in the last t seconds.
func runNamedGlobalTask(name string, t time.Duration, f func() error) bool {
	key := "/@" + name

	jm := JobMarker{
		Node:    serverId,
		Started: time.Now(),
		Type:    "job",
	}

	err := couchbase.Do(key, func(mc *memcached.Client, vb uint16) error {
		resp, err := mc.Add(vb, key, 0, int(t.Seconds()),
			mustEncode(&jm))
		if err != nil {
			return err
		}
		if resp.Status != gomemcached.SUCCESS {
			return fmt.Errorf("Wanted success, got %v", resp.Status)
		}
		return nil
	})

	if err == nil {
		err = f()
		if err != nil {
			log.Printf("Error running periodic task %#v: %v", name, err)
		}
		return true
	}

	return false
}

func heartbeat() {
	for {
		u, err := url.Parse(*couchbaseServer)
		c, err := net.Dial("tcp", u.Host)
		localAddr := ""
		if err == nil {
			localAddr = strings.Split(c.LocalAddr().String(), ":")[0]
			c.Close()
		}

		aboutMe := StorageNode{
			Addr:     localAddr,
			Type:     "node",
			Time:     time.Now().UTC(),
			BindAddr: *bindAddr,
			Hash:     *hashType,
		}

		err = couchbase.Set("/"+serverId, aboutMe)
		if err != nil {
			log.Printf("Failed to record a heartbeat: %v", err)
		}
		time.Sleep(*heartFreq)
	}
}

func verifyObjectHash(h string) error {
	fn := hashFilename(*root, h)
	f, err := os.Open(fn)
	if err != nil {
		return err
	}
	defer f.Close()

	sh := getHash()
	_, err = io.Copy(sh, f)
	if err != nil {
		return err
	}

	hstring := hex.EncodeToString(sh.Sum([]byte{}))
	if h != hstring {
		err = os.Remove(fn)
		if err != nil {
			log.Printf("Error removing corrupt file %v: %v", err)
		}
		return fmt.Errorf("Hash from disk of %v was %v", h, hstring)
	}
	return nil
}

func verifyWorker(ch chan os.FileInfo) {
	for info := range ch {
		err := verifyObjectHash(info.Name())
		if err == nil {
			recordBlobOwnership(info.Name(), info.Size())
		} else {
			log.Printf("Invalid hash for object %v found at verification: %v",
				info.Name(), err)
			removeBlobOwnershipRecord(info.Name(), serverId)
		}
	}
}

func reconcile() error {
	explen := getHash().Size() * 2

	vch := make(chan os.FileInfo)
	defer close(vch)

	for i := 0; i < *verifyWorkers; i++ {
		go verifyWorker(vch)
	}

	return filepath.Walk(*root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && !strings.HasPrefix(info.Name(), "tmp") &&
			len(info.Name()) == explen {

			vch <- info

			return err
		}
		return nil
	})
}

func reconcileLoop() {
	if *reconcileFreq == 0 {
		return
	}
	for {
		err := reconcile()
		if err != nil {
			log.Printf("Error in reconciliation loop: %v", err)
		}
		grabSomeData()
		time.Sleep(*reconcileFreq)
	}
}

func removeBlobOwnershipRecord(h, node string) {
	log.Printf("Cleaning up %v from %v", h, node)

	k := "/" + h
	err := couchbase.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.CAS(vb, k, func(in []byte) ([]byte, memcached.CasOp) {
			ownership := BlobOwnership{}
			err := json.Unmarshal(in, &ownership)
			if err == nil {
				if _, ok := ownership.Nodes[node]; !ok {
					// Skip it fast if we don't have it.
					return nil, memcached.CASQuit
				}
				delete(ownership.Nodes, node)
			} else {
				log.Printf("Error unmarhaling blob removal from %s: %v",
					in, err)
				return nil, memcached.CASQuit
			}

			var rv []byte
			op := memcached.CASStore

			if len(ownership.Nodes) == 0 {
				op = memcached.CASDelete
			} else {
				rv = mustEncode(&ownership)
			}

			return rv, op
		}, 0)
		return err
	})
	if err != nil && err != memcached.CASQuit {
		log.Printf("Error cleaning %v from %v: %v", node, h, err)
	}
}

func cleanupNode(node string) {
	log.Printf("Cleaning up node %v", node)
	vres, err := couchbase.View("cbfs", "node_blobs",
		map[string]interface{}{
			"key":    `"` + node + `"`,
			"limit":  *nodeCleanCount,
			"reduce": false,
			"stale":  false,
		})
	if err != nil {
		log.Printf("Error executing node_blobs view: %v", err)
		return
	}
	foundRows := 0
	for _, r := range vres.Rows {
		removeBlobOwnershipRecord(r.ID[1:], node)
		foundRows++
	}
	if foundRows == 0 && len(vres.Errors) == 0 {
		log.Printf("Removing node record: %v", node)
		err = couchbase.Delete("/" + node)
		if err != nil {
			log.Printf("Error deleting %v node record: %v", node, err)
		}
		err = couchbase.Delete("/" + node + "/r")
		if err != nil {
			log.Printf("Error deleting %v node counter: %v", node, err)
		}
	}
}

func checkStaleNodes() error {
	log.Printf("Checking stale nodes")
	nl, err := findAllNodes()
	if err != nil {
		return err
	}

	for _, node := range nl {
		d := time.Since(node.Time)

		if d > *staleNodeLimit {
			if node.name == serverId {
				log.Printf("Would've cleaned up myself after %v",
					d)
				continue
			}
			log.Printf("  Node %v missed heartbeat schedule: %v",
				node.name, d)
			go cleanupNode(node.name)
		} else {
			log.Printf("%v is ok at %v", node.name, d)
		}
	}
	return nil
}

func garbageCollectBlobs() error {
	log.Printf("Garbage collecting blobs without any file references")

	viewRes := struct {
		Rows []struct {
			Key []string
		}
	}{}

	// we hit this view descending because we want file sorted before blob
	// the fact that we walk the list backwards hopefully not too awkward
	err := couchbase.ViewCustom("cbfs", "file_blobs",
		map[string]interface{}{
			"stale": false, "descending": true}, &viewRes)

	if err != nil {
		return err
	}

	lastBlob := ""
	count := 0
	for _, r := range viewRes.Rows {
		blobId := r.Key[0]
		typeFlag := r.Key[1]
		blobNode := r.Key[2]

		switch typeFlag {
		case "file":
			lastBlob = blobId
		case "blob":
			if blobId != lastBlob {
				go garbageCollectBlobFromNode(blobId, blobNode)
				count++
			}
		}

	}
	log.Printf("Scheduled %d blobs for deletion", count)
	return nil
}

func garbageCollectBlobFromNode(oid, sid string) {
	if sid == serverId {
		//local delete
		err := os.Remove(hashFilename(*root, oid))
		if err != nil {
			log.Printf("Error removing blob, already deleted? %v", err)
		}
	} else {
		//remote
		remote := StorageNode{}
		remotekey := "/" + sid
		err := couchbase.Get(remotekey, &remote)
		if err != nil {
			// will state node cleanup these records or should i?
			log.Printf("No record of this node")
			return
		}

		req, err := http.NewRequest("DELETE", remote.BlobURL(oid), nil)
		resp, err := http.DefaultClient.Do(req)

		if err != nil {
			log.Printf("Error deleting oid %s from node %s, %v",
				oid, sid, err)
			return
		}

		// non-obvious to me at first, but also with 404 we
		// should also remove the blob ownership
		if resp.StatusCode != 204 && resp.StatusCode != 404 {
			log.Printf("Unexpected status %v deleting %s from node %s",
				resp.Status, oid, sid)
			return
		}
	}
	removeBlobOwnershipRecord(oid, sid)
	log.Printf("Removed blob: %v from node %v", oid, sid)
}

type fetchSpec struct {
	oid  string
	node string
}

func dataInitFetchOne(h, u string) error {
	f, err := NewHashRecord(*root, h)
	if err != nil {
		return err
	}
	defer f.Close()

	resp, err := http.Get(u)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	h, l, err := f.Process(resp.Body)
	if err != nil {
		return err
	}
	return recordBlobOwnership(h, l)
}

func dataInitFetcher(nm map[string]StorageNode, ch <-chan fetchSpec) {
	for fs := range ch {
		node, found := nm[fs.node]
		if !found {
			log.Printf("couldn't find %v", fs.node)
			continue
		}
		log.Printf("Fetching %v from %v", fs.oid, node.BlobURL(fs.oid))
		err := dataInitFetchOne(fs.oid, node.BlobURL(fs.oid))
		if err != nil {
			log.Printf("Error fetching %v: %v", fs.oid, err)
		}
	}
}

func grabSomeData() {
	viewRes := struct {
		Rows []struct {
			Id  string
			Doc struct {
				Json struct {
					Nodes map[string]string
				}
			}
		}
	}{}

	// Find some less replicated docs to suck in.
	err := couchbase.ViewCustom("cbfs", "repcounts",
		map[string]interface{}{
			"reduce":       false,
			"include_docs": true,
			"limit":        *maxStartupObjects,
			"startkey":     1,
			"endkey":       1 + *maxStartupRepls,
		},
		&viewRes)

	if err != nil {
		log.Printf("Error finding docs to suck: %v", err)
		return
	}

	nl, err := findRemoteNodes()
	if err != nil {
		log.Printf("Error finding nodes: %v", err)
		return
	}
	nm := map[string]StorageNode{}

	for _, n := range nl {
		nm[n.name] = n
	}

	ch := make(chan fetchSpec, 1000)
	defer close(ch)

	for i := 0; i < 4; i++ {
		go dataInitFetcher(nm, ch)
	}

	for _, r := range viewRes.Rows {
		if _, ok := r.Doc.Json.Nodes[serverId]; !ok {
			for n := range r.Doc.Json.Nodes {
				if n != serverId {
					ch <- fetchSpec{r.Id[1:], n}
				}
			}
		}
	}
}

func runPeriodicJob(name string, job *PeriodicJob) {
	time.Sleep(time.Second * time.Duration(5+rand.Intn(60)))
	for {
		if runNamedGlobalTask(name, job.period, job.f) {
			log.Printf("Attempted job %v", name)
		} else {
			log.Printf("Didn't run job %v", name)
		}
		time.Sleep(job.period + time.Second)
	}
}

func runPeriodicJobs() {
	for n, j := range periodicJobs {
		go runPeriodicJob(n, j)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"strings"

	"github.com/unixvoid/nsproxy/pkg/nslog"
	"gopkg.in/redis.v3"
)

func dnsHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	r.ParseForm()

	dnsType := strings.ToLower(strings.TrimSpace(r.FormValue("dnstype")))
	domain := strings.TrimSpace(r.FormValue("domain"))
	domainValue := strings.TrimSpace(r.FormValue("value"))
	if len(dnsType) == 0 {
		// default to aname entry
		dnsType = "a"
	}

	if dnsType == "cname" {
		// if we are dealing with a CNAME entry fully qualify it
		if string(domainValue[len(domainValue)-1]) != "." {
			domainValue = fmt.Sprintf("%s.", domainValue)
		}
	}

	// make sure domain and value are set
	if (len(domain) == 0) || (len(domainValue) == 0) {
		nslog.Debug.Println("domain or value not set, exiting..")
		w.WriteHeader(http.StatusBadRequest)
	} else {
		// fully qualify the domain name if it is not already:
		if string(domain[len(domain)-1]) != "." {
			domain = fmt.Sprintf("%s.", domain)
		}

		nslog.Debug.Printf("adding domain entry: dns:%s:%s :: %s", dnsType, domain, domainValue)

		// add dns entry dns:<dns_type>:<domain> <domain_value>
		redisClient.Set(fmt.Sprintf("dns:%s:%s", dnsType, domain), domainValue, 0).Err()

		// add dns entry to the list of custom dns names
		redisClient.SAdd("index:dns", fmt.Sprintf("%s:%s", dnsType, domain))

		// return confirmation header to client
		w.Header().Set("x-register", "registered")
		w.WriteHeader(http.StatusOK)
	}
}

func dnsRmHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	r.ParseForm()

	rmType := strings.ToLower(strings.TrimSpace(r.FormValue("dnstype")))
	rmDomain := strings.TrimSpace(r.FormValue("domain"))

	if len(rmDomain) == 0 {
		nslog.Debug.Println("domain not set, exiting..")
		w.WriteHeader(http.StatusBadRequest)
	}

	// fully qualify domain if not done already
	if string(rmDomain[len(rmDomain)-1]) != "." {
		rmDomain = fmt.Sprintf("%s.", rmDomain)
	}

	if len(rmType) == 0 {
		// if type not set, nix them all
		nslog.Debug.Printf("removing all dns types for %s", rmDomain)
		redisClient.Del(fmt.Sprintf("dns:a:%s", rmDomain))
		redisClient.Del(fmt.Sprintf("dns:aaaa:%s", rmDomain))
		redisClient.Del(fmt.Sprintf("dns:cname:%s", rmDomain))

		// remove dns entry from the custom list
		redisClient.SRem("index:dns", fmt.Sprintf("a:%s", rmDomain))
		redisClient.SRem("index:dns", fmt.Sprintf("aaaa:%s", rmDomain))
		redisClient.SRem("index:dns", fmt.Sprintf("cname:%s", rmDomain))
	} else {
		// just remove the specific type
		nslog.Debug.Printf("removing %s entry for %s", rmType, rmDomain)
		redisClient.Del(fmt.Sprintf("dns:%s:%s", rmType, rmDomain))

		// remove dns entry from the custom list
		redisClient.SRem("index:dns", fmt.Sprintf("%s:%s", rmType, rmDomain))
	}
	w.WriteHeader(http.StatusOK)
}

func apiHostsHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	hosts, err := redisClient.SInter("index:live").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	fmt.Fprintln(w, hosts)
}

func dnsHostsHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	hosts, err := redisClient.SInter("index:dns").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	fmt.Fprintln(w, hosts)
}

func apiClustersHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	hosts, err := redisClient.SInter("index:live").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	for _, i := range hosts {
		// we now break at ':' and save the clusters piece
		s := strings.SplitN(i, ":", 2)
		// toss them all into a tmp redis set
		redisClient.SAdd("tmp:cluster:index", s[0])
	}
	// grab the set and delete
	clusters, err := redisClient.SInter("tmp:cluster:index").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	redisClient.Del("tmp:cluster:index")
	fmt.Fprintln(w, clusters)
}

func apiClusterSpecHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	queryCluster := strings.TrimSpace(r.FormValue("cluster"))
	hosts, err := redisClient.SInter("index:live").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}

	for _, i := range hosts {
		// we now break at ':' and save the clusters piece
		s := strings.SplitN(i, ":", 2)
		if s[0] == queryCluster {
			// if the custer matches the query, throw the host in a tmp set
			redisClient.SAdd("tmp:cluster:index", s[1])
		}
	}
	// grab the set and delete
	clusters, err := redisClient.SInter("tmp:cluster:index").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	redisClient.Del("tmp:cluster:index")
	fmt.Fprintln(w, clusters)
}

func apiHostSpecHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	cluster := strings.TrimSpace(r.FormValue("cluster"))
	host := strings.TrimSpace(r.FormValue("host"))

	ip, err := redisClient.Get(fmt.Sprintf("cluster:%s:%s", cluster, host)).Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	fmt.Fprintln(w, ip)
}

func apiDnsSpecHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	dnsType := strings.TrimSpace(r.FormValue("dnstype"))
	domainValue := strings.TrimSpace(r.FormValue("domain"))

	// fully qualify if not already
	if string(domainValue[len(domainValue)-1]) != "." {
		domainValue = fmt.Sprintf("%s.", domainValue)
	}

	// default to a record
	if len(dnsType) == 0 {
		dnsType = "a"
	}

	ip, err := redisClient.Get(fmt.Sprintf("dns:%s:%s", dnsType, domainValue)).Result()
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
	}

	fmt.Fprintln(w, ip)
}
<commit_msg>Fix bug where cname is set but invalid<commit_after>package main

import (
	"fmt"
	"net/http"
	"strings"

	"github.com/unixvoid/nsproxy/pkg/nslog"
	"gopkg.in/redis.v3"
)

func dnsHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	r.ParseForm()

	dnsType := strings.ToLower(strings.TrimSpace(r.FormValue("dnstype")))
	domain := strings.TrimSpace(r.FormValue("domain"))
	domainValue := strings.TrimSpace(r.FormValue("value"))

	// make sure domain and value are set
	if (len(domain) == 0) || (len(domainValue) == 0) {
		nslog.Debug.Println("domain or value not set, exiting..")
		w.WriteHeader(http.StatusBadRequest)
	} else {
		if len(dnsType) == 0 {
			// default to aname entry
			dnsType = "a"
		} else {
			// if dnstype is set, make sure it is something we support
			switch dnsType {
			case
				"a",
				"aaaa",
				"cname":
				break
			default:
				nslog.Debug.Printf("unsupported dnstype '%s', exiting..\n", dnsType)
				w.WriteHeader(http.StatusBadRequest)
				return
			}
		}
		if dnsType == "cname" {
			// if we are dealing with a CNAME entry fully qualify it
			if string(domainValue[len(domainValue)-1]) != "." {
				domainValue = fmt.Sprintf("%s.", domainValue)
			}
		}
		// fully qualify the domain name if it is not already:
		if string(domain[len(domain)-1]) != "." {
			domain = fmt.Sprintf("%s.", domain)
		}

		nslog.Debug.Printf("adding domain entry: dns:%s:%s :: %s", dnsType, domain, domainValue)

		// add dns entry dns:<dns_type>:<domain> <domain_value>
		redisClient.Set(fmt.Sprintf("dns:%s:%s", dnsType, domain), domainValue, 0).Err()

		// add dns entry to the list of custom dns names
		redisClient.SAdd("index:dns", fmt.Sprintf("%s:%s", dnsType, domain))

		// return confirmation header to client
		w.Header().Set("x-register", "registered")
		w.WriteHeader(http.StatusOK)
	}
}

func dnsRmHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	r.ParseForm()

	rmType := strings.ToLower(strings.TrimSpace(r.FormValue("dnstype")))
	rmDomain := strings.TrimSpace(r.FormValue("domain"))

	if len(rmDomain) == 0 {
		nslog.Debug.Println("domain not set, exiting..")
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	// fully qualify domain if not done already
	if string(rmDomain[len(rmDomain)-1]) != "." {
		rmDomain = fmt.Sprintf("%s.", rmDomain)
	}

	if len(rmType) == 0 {
		// if type not set, nix them all
		nslog.Debug.Printf("removing all dns types for %s", rmDomain)
		redisClient.Del(fmt.Sprintf("dns:a:%s", rmDomain))
		redisClient.Del(fmt.Sprintf("dns:aaaa:%s", rmDomain))
		redisClient.Del(fmt.Sprintf("dns:cname:%s", rmDomain))

		// remove dns entry from the custom list
		redisClient.SRem("index:dns", fmt.Sprintf("a:%s", rmDomain))
		redisClient.SRem("index:dns", fmt.Sprintf("aaaa:%s", rmDomain))
		redisClient.SRem("index:dns", fmt.Sprintf("cname:%s", rmDomain))
	} else {
		// just remove the specific type
		nslog.Debug.Printf("removing %s entry for %s", rmType, rmDomain)
		redisClient.Del(fmt.Sprintf("dns:%s:%s", rmType, rmDomain))

		// remove dns entry from the custom list
		redisClient.SRem("index:dns", fmt.Sprintf("%s:%s", rmType, rmDomain))
	}
	w.WriteHeader(http.StatusOK)
}

func apiHostsHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	hosts, err := redisClient.SInter("index:live").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	fmt.Fprintln(w, hosts)
}

func dnsHostsHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	hosts, err := redisClient.SInter("index:dns").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	fmt.Fprintln(w, hosts)
}

func apiClustersHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	hosts, err := redisClient.SInter("index:live").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	for _, i := range hosts {
		// we now break at ':' and save the clusters piece
		s := strings.SplitN(i, ":", 2)
		// toss them all into a tmp redis set
		redisClient.SAdd("tmp:cluster:index", s[0])
	}
	// grab the set and delete
	clusters, err := redisClient.SInter("tmp:cluster:index").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	redisClient.Del("tmp:cluster:index")
	fmt.Fprintln(w, clusters)
}

func apiClusterSpecHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	queryCluster := strings.TrimSpace(r.FormValue("cluster"))
	hosts, err := redisClient.SInter("index:live").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}

	for _, i := range hosts {
		// we now break at ':' and save the clusters piece
		s := strings.SplitN(i, ":", 2)
		if s[0] == queryCluster {
			// if the custer matches the query, throw the host in a tmp set
			redisClient.SAdd("tmp:cluster:index", s[1])
		}
	}
	// grab the set and delete
	clusters, err := redisClient.SInter("tmp:cluster:index").Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	redisClient.Del("tmp:cluster:index")
	fmt.Fprintln(w, clusters)
}

func apiHostSpecHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	cluster := strings.TrimSpace(r.FormValue("cluster"))
	host := strings.TrimSpace(r.FormValue("host"))

	ip, err := redisClient.Get(fmt.Sprintf("cluster:%s:%s", cluster, host)).Result()
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
	}
	fmt.Fprintln(w, ip)
}

func apiDnsSpecHandler(w http.ResponseWriter, r *http.Request, redisClient *redis.Client) {
	dnsType := strings.TrimSpace(r.FormValue("dnstype"))
	domainValue := strings.TrimSpace(r.FormValue("domain"))

	// fully qualify if not already
	if string(domainValue[len(domainValue)-1]) != "." {
		domainValue = fmt.Sprintf("%s.", domainValue)
	}

	// default to a record
	if len(dnsType) == 0 {
		dnsType = "a"
	}

	ip, err := redisClient.Get(fmt.Sprintf("dns:%s:%s", dnsType, domainValue)).Result()
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
	}

	fmt.Fprintln(w, ip)
}
<|endoftext|>