<commit_before>// Golang URL Minifier
// Creates tiny URLS from a given URL.
// https://github.com/nickvellios/Golang-URL-Minifier
// Nick Vellios
// 11/23/2016

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	_ "github.com/lib/pq"
	"html/template"
	"math"
	"net"
	"net/http"
	"net/url"
	"strings"
)

var templates = template.Must(template.ParseFiles(
	templateDir+"index.html",
	templateDir+"header.html",
	templateDir+"footer.html"))

type Tiny struct {
	URL       string
	Path      string
	IP        string
	Timestamp string
	ID        string
}

type URLResponse struct {
	URL   string `json:"url"`
	Error string `json:"error"`
}

// There may be URLs longer, but to avoid attacks we don't want them.
const (
	MAX_URL = 1024 // Max length of URL to shrink
	MAX_REQ = 10   // Max number of requests per hour
)

// Send a JSON result back to the client with given status code
func writeResponse(w http.ResponseWriter, code int, u, error string) error {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	ur := &URLResponse{u, error}
	err := json.NewEncoder(w).Encode(ur)

	return err
}

func renderTemplate(w http.ResponseWriter, tmpl string, data []map[string]string) {
	err := templates.ExecuteTemplate(w, tmpl, data)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// HTTP handler for /generate/ which is the API.  Takes a long URL and generates a tiny URL.
func (udb *urlDB) generateHandler(w http.ResponseWriter, r *http.Request) {
	// In order to support MaxBytesReader and r.ContentLength I am limiting to POST/PUT requests
	switch r.Method {
	case "GET":
		writeResponse(w, 405, "", "API only supports POST requests.")
		return
	}

	// Make sure we aren't flooded with excess data from some asshole
	if r.ContentLength > MAX_URL {
		writeResponse(w, 413, "", "URL exceeds maximum length of 1024 characters")
		return
	}

	// Check again to be sure we aren't flooded with excess data
	r.Body = http.MaxBytesReader(w, r.Body, MAX_URL)
	err := r.ParseForm()
	if err != nil {
		writeResponse(w, 413, "", "URL exceeds maximum length of 1024 characters")
		return
	}

	urlf := r.FormValue("url")

	up, err := url.Parse(urlf)
	if err != nil {
		writeResponse(w, 500, "", "Internal Server Error")
		return
	}

	// We expect a URL has at least one period.
	if !strings.Contains(up.Host, ".") {
		writeResponse(w, 400, "", "Invalid URL")
		return
	}
	// 301/302 redirects fail without a valid URL.  Our site frondend checks for this and adds a
	// http:// prefix if no URL scheme is found but we will check and do the same for API requests
	if up.Scheme != "http" && up.Scheme != "https" && up.Scheme != "ftp" {
		urlf = "http://" + urlf
	}

	// To prevent someone from building predictive redirect chains to try and overload us
	if strings.Contains(urlf, baseURL) {
		writeResponse(w, 400, "", "Short urls pointing to " + baseURL + " are not allowed")
		return
	}

	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		writeResponse(w, 500, "", "Internal Server Error") // User IP is not IP:port.  Something is fishy, kill it just in case.
		return
	}

	tiny := &Tiny{URL: urlf, IP: ip}

	// Limit to MAX_REQ requests per hour per IP
	if tiny.throttleCheck(udb.db) {
		tiny.save(udb.db)
		writeResponse(w, 200, baseURL+tiny.Path, "")
	} else {
		writeResponse(w, 429, "", "You're doing that too often.  Slow down")
	}
}

// HTTP handler for / path
func (udb *urlDB) rootHandler(w http.ResponseWriter, r *http.Request) {
	urlf := r.URL.Path[len("/"):]

	if len(urlf) > 0 && urlf != "favicon.ico" {
		t := &Tiny{Path: urlf}
		t.load(udb.db)
		fmt.Println("Redirecting to: ", t.URL)
		http.Redirect(w, r, t.URL, 302)

		return
	}

	renderTemplate(w, "index", nil)
}

// Save URL to DB, get unique ID, generate tiny path from the ID, update the DB.
func (t *Tiny) save(db *sql.DB) int {
	var lastInsertId int
	err := db.QueryRow("INSERT INTO url_map(path, url, ip) VALUES($1,$2,$3) returning id;", "", t.URL, t.IP).Scan(&lastInsertId)
	checkDBErr(err)

	stmt, err := db.Prepare("UPDATE url_map SET path=$1 WHERE id=$2")
	checkDBErr(err)
	path := generateCode(lastInsertId)
	t.Path = path
	_, err = stmt.Exec(t.Path, lastInsertId)
	checkDBErr(err)

	return lastInsertId
}

// Load URL from DB.
func (t *Tiny) load(db *sql.DB) {
	rows, err := db.Query("SELECT url FROM url_map WHERE path = $1", t.Path)
	checkDBErr(err)

	for rows.Next() {
		err := rows.Scan(&t.URL)
		checkDBErr(err)
	}
}

// Check if user has used service more than MAX_REQ times in an hour.
func (t *Tiny) throttleCheck(db *sql.DB) bool {
	rows, err := db.Query("SELECT COUNT(*) as count FROM url_map WHERE ip = $1 AND t_stamp > CURRENT_TIMESTAMP - INTERVAL '1 hour'", t.IP)
	checkDBErr(err)

	var count int
	for rows.Next() {
		err := rows.Scan(&count)
		checkDBErr(err)
	}

	return count < MAX_REQ
}

// Generate a unique A-Z/0-9 based URL from a given number input.
func generateCode(number int) string {
	var out []byte
	codes := []byte("abcdefghjkmnpqrstuvwxyz23456789ABCDEFGHJKMNPQRSTUVWXYZ")

	for number > 53 {
		key := number % 54
		number = int(math.Floor(float64(number)/54) - 1)
		out = append(out, []byte(codes[key : key+1])[0])
	}

	return string(append(out, codes[number]))
}
<commit_msg>gofmt cleanup<commit_after>// Golang URL Minifier
// Creates tiny URLS from a given URL.
// https://github.com/nickvellios/Golang-URL-Minifier
// Nick Vellios
// 11/23/2016

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	_ "github.com/lib/pq"
	"html/template"
	"math"
	"net"
	"net/http"
	"net/url"
	"strings"
)

var templates = template.Must(template.ParseFiles(
	templateDir+"index.html",
	templateDir+"header.html",
	templateDir+"footer.html"))

type Tiny struct {
	URL       string
	Path      string
	IP        string
	Timestamp string
	ID        string
}

type URLResponse struct {
	URL   string `json:"url"`
	Error string `json:"error"`
}

// There may be URLs longer, but to avoid attacks we don't want them.
const (
	MAX_URL = 1024 // Max length of URL to shrink
	MAX_REQ = 10   // Max number of requests per hour
)

// Send a JSON result back to the client with given status code
func writeResponse(w http.ResponseWriter, code int, u, error string) error {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	ur := &URLResponse{u, error}
	err := json.NewEncoder(w).Encode(ur)

	return err
}

func renderTemplate(w http.ResponseWriter, tmpl string, data []map[string]string) {
	err := templates.ExecuteTemplate(w, tmpl, data)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// HTTP handler for /generate/ which is the API.  Takes a long URL and generates a tiny URL.
func (udb *urlDB) generateHandler(w http.ResponseWriter, r *http.Request) {
	// In order to support MaxBytesReader and r.ContentLength I am limiting to POST/PUT requests
	switch r.Method {
	case "GET":
		writeResponse(w, 405, "", "API only supports POST requests.")
		return
	}

	// Make sure we aren't flooded with excess data from some asshole
	if r.ContentLength > MAX_URL {
		writeResponse(w, 413, "", "URL exceeds maximum length of 1024 characters")
		return
	}

	// Check again to be sure we aren't flooded with excess data
	r.Body = http.MaxBytesReader(w, r.Body, MAX_URL)
	err := r.ParseForm()
	if err != nil {
		writeResponse(w, 413, "", "URL exceeds maximum length of 1024 characters")
		return
	}

	urlf := r.FormValue("url")

	up, err := url.Parse(urlf)
	if err != nil {
		writeResponse(w, 500, "", "Internal Server Error")
		return
	}

	// We expect a URL has at least one period.
	if !strings.Contains(up.Host, ".") {
		writeResponse(w, 400, "", "Invalid URL")
		return
	}
	// 301/302 redirects fail without a valid URL.  Our site frondend checks for this and adds a
	// http:// prefix if no URL scheme is found but we will check and do the same for API requests
	if up.Scheme != "http" && up.Scheme != "https" && up.Scheme != "ftp" {
		urlf = "http://" + urlf
	}

	// To prevent someone from building predictive redirect chains to try and overload us
	if strings.Contains(urlf, baseURL) {
		writeResponse(w, 400, "", "Short urls pointing to "+baseURL+" are not allowed")
		return
	}

	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		writeResponse(w, 500, "", "Internal Server Error") // User IP is not IP:port.  Something is fishy, kill it just in case.
		return
	}

	tiny := &Tiny{URL: urlf, IP: ip}

	// Limit to MAX_REQ requests per hour per IP
	if tiny.throttleCheck(udb.db) {
		tiny.save(udb.db)
		writeResponse(w, 200, baseURL+tiny.Path, "")
	} else {
		writeResponse(w, 429, "", "You're doing that too often.  Slow down")
	}
}

// HTTP handler for / path
func (udb *urlDB) rootHandler(w http.ResponseWriter, r *http.Request) {
	urlf := r.URL.Path[len("/"):]

	if len(urlf) > 0 && urlf != "favicon.ico" {
		t := &Tiny{Path: urlf}
		t.load(udb.db)
		fmt.Println("Redirecting to: ", t.URL)
		http.Redirect(w, r, t.URL, 302)

		return
	}

	renderTemplate(w, "index", nil)
}

// Save URL to DB, get unique ID, generate tiny path from the ID, update the DB.
func (t *Tiny) save(db *sql.DB) int {
	var lastInsertId int
	err := db.QueryRow("INSERT INTO url_map(path, url, ip) VALUES($1,$2,$3) returning id;", "", t.URL, t.IP).Scan(&lastInsertId)
	checkDBErr(err)

	stmt, err := db.Prepare("UPDATE url_map SET path=$1 WHERE id=$2")
	checkDBErr(err)
	path := generateCode(lastInsertId)
	t.Path = path
	_, err = stmt.Exec(t.Path, lastInsertId)
	checkDBErr(err)

	return lastInsertId
}

// Load URL from DB.
func (t *Tiny) load(db *sql.DB) {
	rows, err := db.Query("SELECT url FROM url_map WHERE path = $1", t.Path)
	checkDBErr(err)

	for rows.Next() {
		err := rows.Scan(&t.URL)
		checkDBErr(err)
	}
}

// Check if user has used service more than MAX_REQ times in an hour.
func (t *Tiny) throttleCheck(db *sql.DB) bool {
	rows, err := db.Query("SELECT COUNT(*) as count FROM url_map WHERE ip = $1 AND t_stamp > CURRENT_TIMESTAMP - INTERVAL '1 hour'", t.IP)
	checkDBErr(err)

	var count int
	for rows.Next() {
		err := rows.Scan(&count)
		checkDBErr(err)
	}

	return count < MAX_REQ
}

// Generate a unique A-Z/0-9 based URL from a given number input.
func generateCode(number int) string {
	var out []byte
	codes := []byte("abcdefghjkmnpqrstuvwxyz23456789ABCDEFGHJKMNPQRSTUVWXYZ")

	for number > 53 {
		key := number % 54
		number = int(math.Floor(float64(number)/54) - 1)
		out = append(out, []byte(codes[key : key+1])[0])
	}

	return string(append(out, codes[number]))
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/dfordsoft/golib/ic"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"
)

func init() {
	registerNovelSiteHandler(&NovelSiteHandler{
		Match:    isPiaotian,
		Download: dlPiaotian,
	})
}

func isPiaotian(u string) bool {
	r, _ := regexp.Compile(`http://www\.piaotian\.com/html/[0-9]/[0-9]+/`)
	if r.MatchString(u) {
		return true
	}
	r, _ = regexp.Compile(`http://www\.piaotian\.com/bookinfo/[0-9]/[0-9]+\.html`)
	if r.MatchString(u) {
		return true
	}
	return false
}

func dlPiaotianPage(u string) (c []byte) {
	client := &http.Client{
		Timeout: 60 * time.Second,
	}
	retry := 0
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		log.Println("piaotian - Could not parse novel page request:", err)
		return
	}

	req.Header.Set("Referer", "http://www.piaotian.com/")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:45.0) Gecko/20100101 Firefox/45.0")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Header.Set("accept-language", `en-US,en;q=0.8`)
	req.Header.Set("Upgrade-Insecure-Requests", "1")
doRequest:
	resp, err := client.Do(req)
	if err != nil {
		log.Println("piaotian - Could not send novel page request:", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("piaotian - novel page request not 200")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	c, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("piaotian - novel page content reading failed")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}
	c = ic.Convert("gbk", "utf-8", c)
	c = bytes.Replace(c, []byte("\r\n"), []byte(""), -1)
	c = bytes.Replace(c, []byte("\r"), []byte(""), -1)
	c = bytes.Replace(c, []byte("\n"), []byte(""), -1)
	idx := bytes.Index(c, []byte("

    "))
	if idx > 1 {
		c = c[idx+17:]
	}
	idx = bytes.Index(c, []byte("
"))
	if idx > 1 {
		c = c[:idx]
	}
	c = bytes.Replace(c, []byte("
    "), []byte(""), -1)
	c = bytes.Replace(c, []byte("    "), []byte(""), -1)
	return
}

func dlPiaotian(u string) {
	tocURL := u
	r, _ := regexp.Compile(`http://www\.piaotian\.com/bookinfo/([0-9])/([0-9]+)\.html`)
	if r.MatchString(u) {
		ss := r.FindAllStringSubmatch(u, -1)
		s := ss[0]
		tocURL = fmt.Sprintf("http://www.piaotian.com/html/%s/%s/", s[1], s[2])
	}
	fmt.Println("download book from", tocURL)

	client := &http.Client{
		Timeout: 60 * time.Second,
	}
	retry := 0
	req, err := http.NewRequest("GET", tocURL, nil)
	if err != nil {
		log.Println("piaotian - Could not parse novel request:", err)
		return
	}

	req.Header.Set("Referer", "http://www.piaotian.com/")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:45.0) Gecko/20100101 Firefox/45.0")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Header.Set("accept-language", `en-US,en;q=0.8`)
	req.Header.Set("Upgrade-Insecure-Requests", "1")
doRequest:
	resp, err := client.Do(req)
	if err != nil {
		log.Println("piaotian - Could not send novel request:", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("piaotian - novel request not 200")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	r, _ = regexp.Compile(`^<li><a\shref="([0-9]+\.html)">([^<]+)</a></li>$`)
	scanner := bufio.NewScanner(resp.Body)
	scanner.Split(bufio.ScanLines)

	contentHTML, err := os.OpenFile(`content.html`, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Println("opening file content.html for writing failed ", err)
		return
	}

	contentHTMLTemplate := `<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.1//EN" "http://www.w3.org/TR/xhtml11/DTD/xhtml11.dtd">

		<meta http-equiv="Content-Type" content="text/html; charset=utf-8">
		<title>Get Novel</title>

Get Novel

	Go straight to first item
	06/17 06:31

目录

		%s

	<mbp:pagebreak></mbp:pagebreak>

		%s

	`

	var toc, content []string
	var navPoint []string
	for scanner.Scan() {
		line := scanner.Text()
		// convert from gbk to UTF-8
		l := ic.ConvertString("gbk", "utf-8", line)
		if r.MatchString(l) {
			ss := r.FindAllStringSubmatch(l, -1)
			s := ss[0]
			finalURL := fmt.Sprintf("%s%s", tocURL, s[1])
			idx := len(toc)
			toc = append(toc, fmt.Sprintf(`<li><a href="#article_%d">%s</a></li>`, idx, s[2]))
			c := dlPiaotianPage(finalURL)
			content = append(content, fmt.Sprintf(`

				  %s

%s

`, idx, finalURL, s[2], string(c)))
			navPoint = append(navPoint, fmt.Sprintf(`
				<navPoint class="chapter" id="%d" playOrder="1">
					<navLabel><text>%s</text></navLabel>
					<content src="content.html#article_%d" />
				</navPoint>
				`, idx, s[2], idx))

			fmt.Println(s[2], finalURL, len(c), "bytes")
		}
	}
	contentHTML.WriteString(fmt.Sprintf(contentHTMLTemplate, strings.Join(toc, "\n"), strings.Join(content, "\n")))
	contentHTML.Close()

	tocNCX, err := os.OpenFile("toc.ncx", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Println("opening file toc.ncx for writing failed ", err)
		return
	}

	tocNCXTemplate := `<?xml version="1.0" encoding="UTF-8"?>
	<ncx xmlns="http://www.daisy.org/z3986/2005/ncx/" version="2005-1" xml:lang="zh-CN">

	<meta name="dtb:uid" content="11562530804848545888" />
	<meta name="dtb:depth" content="4" />
	<meta name="dtb:totalPageCount" content="0" />
	<meta name="dtb:maxPageNumber" content="0" />

	<docTitle><text>Get Novel</text></docTitle>
	<docAuthor><text>类库</text></docAuthor>
	<navMap>
		<navPoint class="book">
			<navLabel><text>Get Novel</text></navLabel>
			<content src="content.html" />
			%s
		</navPoint>
	</navMap>
	</ncx>`

	tocNCX.WriteString(fmt.Sprintf(tocNCXTemplate, strings.Join(navPoint, "\n")))
	tocNCX.Close()

	contentOPF, err := os.OpenFile("content.opf", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Println("opening file content.opf for writing failed ", err)
		return
	}
	contentOPFTemplate := `<?xml version="1.0" encoding="utf-8"?>
	<package xmlns="http://www.idpf.org/2007/opf" version="2.0" unique-identifier="uid">
	<metadata>
	<dc-metadata xmlns:dc="http://purl.org/dc/elements/1.1/" xmlns:opf="http://www.idpf.org/2007/opf">
		<dc:title>Get Novel</dc:title>
		<dc:language>zh-CN</dc:language>
		<dc:identifier id="uid">115625308048485458882013-06-16T22:31:08Z</dc:identifier>
		<dc:creator>kindlereader</dc:creator>
		<dc:publisher>kindlereader</dc:publisher>
		<dc:subject>Get Novel</dc:subject>
		<dc:date>2013-06-16T22:31:08Z</dc:date>
		<dc:description></dc:description>
	</dc-metadata>

	</metadata>
	<manifest>
		<item id="content" media-type="application/xhtml+xml" href="content.html"></item>
		<item id="toc" media-type="application/x-dtbncx+xml" href="toc.ncx"></item>
	</manifest>

	<spine toc="toc">
		<itemref idref="content"/>
	</spine>

	<guide>
		<reference type="start" title="start" href="content.html#content"></reference>
		<reference type="toc" title="toc" href="content.html#toc"></reference>
		<reference type="text" title="cover" href="content.html#cover"></reference>
	</guide>
	</package>
	`
	contentOPF.WriteString(contentOPFTemplate)
	contentOPF.Close()
}
<commit_msg>(*)fix title<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/dfordsoft/golib/ic"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"
)

func init() {
	registerNovelSiteHandler(&NovelSiteHandler{
		Match:    isPiaotian,
		Download: dlPiaotian,
	})
}

func isPiaotian(u string) bool {
	r, _ := regexp.Compile(`http://www\.piaotian\.com/html/[0-9]/[0-9]+/`)
	if r.MatchString(u) {
		return true
	}
	r, _ = regexp.Compile(`http://www\.piaotian\.com/bookinfo/[0-9]/[0-9]+\.html`)
	if r.MatchString(u) {
		return true
	}
	return false
}

func dlPiaotianPage(u string) (c []byte) {
	client := &http.Client{
		Timeout: 60 * time.Second,
	}
	retry := 0
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		log.Println("piaotian - Could not parse novel page request:", err)
		return
	}

	req.Header.Set("Referer", "http://www.piaotian.com/")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:45.0) Gecko/20100101 Firefox/45.0")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Header.Set("accept-language", `en-US,en;q=0.8`)
	req.Header.Set("Upgrade-Insecure-Requests", "1")
doRequest:
	resp, err := client.Do(req)
	if err != nil {
		log.Println("piaotian - Could not send novel page request:", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("piaotian - novel page request not 200")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	c, err = ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Println("piaotian - novel page content reading failed")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}
	c = ic.Convert("gbk", "utf-8", c)
	c = bytes.Replace(c, []byte("\r\n"), []byte(""), -1)
	c = bytes.Replace(c, []byte("\r"), []byte(""), -1)
	c = bytes.Replace(c, []byte("\n"), []byte(""), -1)
	idx := bytes.Index(c, []byte("

    "))
	if idx > 1 {
		c = c[idx+17:]
	}
	idx = bytes.Index(c, []byte("
"))
	if idx > 1 {
		c = c[:idx]
	}
	c = bytes.Replace(c, []byte("
    "), []byte(""), -1)
	c = bytes.Replace(c, []byte("    "), []byte(""), -1)
	return
}

func dlPiaotian(u string) {
	tocURL := u
	r, _ := regexp.Compile(`http://www\.piaotian\.com/bookinfo/([0-9])/([0-9]+)\.html`)
	if r.MatchString(u) {
		ss := r.FindAllStringSubmatch(u, -1)
		s := ss[0]
		tocURL = fmt.Sprintf("http://www.piaotian.com/html/%s/%s/", s[1], s[2])
	}
	fmt.Println("download book from", tocURL)

	client := &http.Client{
		Timeout: 60 * time.Second,
	}
	retry := 0
	req, err := http.NewRequest("GET", tocURL, nil)
	if err != nil {
		log.Println("piaotian - Could not parse novel request:", err)
		return
	}

	req.Header.Set("Referer", "http://www.piaotian.com/")
	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:45.0) Gecko/20100101 Firefox/45.0")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Header.Set("accept-language", `en-US,en;q=0.8`)
	req.Header.Set("Upgrade-Insecure-Requests", "1")
doRequest:
	resp, err := client.Do(req)
	if err != nil {
		log.Println("piaotian - Could not send novel request:", err)
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		log.Println("piaotian - novel request not 200")
		retry++
		if retry < 3 {
			time.Sleep(3 * time.Second)
			goto doRequest
		}
		return
	}

	scanner := bufio.NewScanner(resp.Body)
	scanner.Split(bufio.ScanLines)

	contentHTML, err := os.OpenFile(`content.html`, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Println("opening file content.html for writing failed ", err)
		return
	}

	contentHTMLTemplate := `<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.1//EN" "http://www.w3.org/TR/xhtml11/DTD/xhtml11.dtd">

		<meta http-equiv="Content-Type" content="text/html; charset=utf-8">
		<title>%s</title>

%s

	跳到第一篇
%s

目录

		%s

	<mbp:pagebreak></mbp:pagebreak>

		%s

	`

	var title string
	var toc, content []string
	var navPoint []string
	r, _ = regexp.Compile(`^<li><a\shref="([0-9]+\.html)">([^<]+)</a></li>$`)
	for scanner.Scan() {
		line := scanner.Text()
		// convert from gbk to UTF-8
		l := ic.ConvertString("gbk", "utf-8", line)
		if title == "" {
			re, _ := regexp.Compile(`^<h1>([^<]+)</h1>$`)
			ss := re.FindAllStringSubmatch(l, -1)
			if len(ss) > 0 && len(ss[0]) > 0 {
				s := ss[0]
				title = s[1]
				idx := strings.Index(title, `最新章节`)
				if idx > 0 {
					title = title[:idx]
				}
				continue
			}
		}
		if r.MatchString(l) {
			ss := r.FindAllStringSubmatch(l, -1)
			s := ss[0]
			finalURL := fmt.Sprintf("%s%s", tocURL, s[1])
			idx := len(toc)
			toc = append(toc, fmt.Sprintf(`<li><a href="#article_%d">%s</a></li>`, idx, s[2]))
			c := dlPiaotianPage(finalURL)
			content = append(content, fmt.Sprintf(`

				  %s

%s

`, idx, finalURL, s[2], string(c)))
			navPoint = append(navPoint, fmt.Sprintf(`
				<navPoint class="chapter" id="%d" playOrder="1">
					<navLabel><text>%s</text></navLabel>
					<content src="content.html#article_%d" />
				</navPoint>
				`, idx, s[2], idx))

			fmt.Println(s[2], finalURL, len(c), "bytes")
		}
	}
	contentHTML.WriteString(fmt.Sprintf(contentHTMLTemplate, title, title, time.Now().String(),
		strings.Join(toc, "\n"), strings.Join(content, "\n")))
	contentHTML.Close()

	tocNCX, err := os.OpenFile("toc.ncx", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Println("opening file toc.ncx for writing failed ", err)
		return
	}

	uid := time.Now().Nanosecond()
	tocNCXTemplate := `<?xml version="1.0" encoding="UTF-8"?>
	<ncx xmlns="http://www.daisy.org/z3986/2005/ncx/" version="2005-1" xml:lang="zh-CN">

	<meta name="dtb:uid" content="%d" />
	<meta name="dtb:depth" content="4" />
	<meta name="dtb:totalPageCount" content="0" />
	<meta name="dtb:maxPageNumber" content="0" />

	<docTitle><text>%s</text></docTitle>
	<docAuthor><text>类库大魔王</text></docAuthor>
	<navMap>
		<navPoint class="book">
			<navLabel><text>%s</text></navLabel>
			<content src="content.html" />
			%s
		</navPoint>
	</navMap>
	</ncx>`

	tocNCX.WriteString(fmt.Sprintf(tocNCXTemplate, uid, title, title, strings.Join(navPoint, "\n")))
	tocNCX.Close()

	contentOPF, err := os.OpenFile("content.opf", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		log.Println("opening file content.opf for writing failed ", err)
		return
	}
	contentOPFTemplate := `<?xml version="1.0" encoding="utf-8"?>
	<package xmlns="http://www.idpf.org/2007/opf" version="2.0" unique-identifier="uid">
	<metadata>
	<dc-metadata xmlns:dc="http://purl.org/dc/elements/1.1/" xmlns:opf="http://www.idpf.org/2007/opf">
		<dc:title>%s</dc:title>
		<dc:language>zh-CN</dc:language>
		<dc:identifier id="uid">%d%s</dc:identifier>
		<dc:creator>GetNovel</dc:creator>
		<dc:publisher>类库大魔王</dc:publisher>
		<dc:subject>%s</dc:subject>
		<dc:date>%s</dc:date>
		<dc:description></dc:description>
	</dc-metadata>

	</metadata>
	<manifest>
		<item id="content" media-type="application/xhtml+xml" href="content.html"></item>
		<item id="toc" media-type="application/x-dtbncx+xml" href="toc.ncx"></item>
	</manifest>

	<spine toc="toc">
		<itemref idref="content"/>
	</spine>

	<guide>
		<reference type="start" title="start" href="content.html#content"></reference>
		<reference type="toc" title="toc" href="content.html#toc"></reference>
		<reference type="text" title="cover" href="content.html#cover"></reference>
	</guide>
	</package>
	`
	contentOPF.WriteString(fmt.Sprintf(contentOPFTemplate, title, uid, time.Now().String(), title, time.Now().String()))
	contentOPF.Close()
}
<|endoftext|>


<commit_before>package frames

import (
	"reflect"
	"testing"
)

func TestPktEncoding(t *testing.T) {
	t.Parallel()
	tests := []struct {
		pkt FramePacket
		exp []byte
	}{
		{FramePacket{Cmd: FrameOpen},
			[]byte{0, 0, 0, 0, 0, 0}},
		{FramePacket{Cmd: FrameClose, Channel: 923},
			[]byte{0, 0, 3, 0x9b, 1, 0}},
		{FramePacket{Cmd: FrameOpen, Status: FrameError, Channel: 13},
			[]byte{0, 0, 0, 13, 0, 1}},
		{FramePacket{Cmd: FrameData, Channel: 11, Data: []byte("hi")},
			[]byte{0, 2, 0, 11, 2, 0, 'h', 'i'}},
	}

	for _, test := range tests {
		got := test.pkt.Bytes()
		if !reflect.DeepEqual(got, test.exp) {
			t.Errorf("Error encoding %v\nExpected:\n%#v\nGot:\n%#v",
				test.pkt, test.exp, got)
		}
		t.Logf("Packet: %v", test.pkt)
	}
}

func benchEncoding(b *testing.B, size int) {
	pkt := FramePacket{
		Cmd:     FrameData,
		Channel: 8184,
		Data:    make([]byte, size),
	}

	b.SetBytes(int64(len(pkt.Bytes())))

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		pkt.Bytes()
	}
}

func BenchmarkEncoding0(b *testing.B) {
	benchEncoding(b, 0)
}

func BenchmarkEncoding8(b *testing.B) {
	benchEncoding(b, 8)
}

func BenchmarkEncoding16(b *testing.B) {
	benchEncoding(b, 16)
}

func BenchmarkEncoding64(b *testing.B) {
	benchEncoding(b, 64)
}

func BenchmarkEncoding256(b *testing.B) {
	benchEncoding(b, 256)
}

func BenchmarkEncoding1024(b *testing.B) {
	benchEncoding(b, 1024)
}

func BenchmarkEncoding8192(b *testing.B) {
	benchEncoding(b, 8192)
}
<commit_msg>Added test for frame errors<commit_after>package frames

import (
	"reflect"
	"testing"
)

func TestPktEncoding(t *testing.T) {
	t.Parallel()
	tests := []struct {
		pkt FramePacket
		exp []byte
	}{
		{FramePacket{Cmd: FrameOpen},
			[]byte{0, 0, 0, 0, 0, 0}},
		{FramePacket{Cmd: FrameClose, Channel: 923},
			[]byte{0, 0, 3, 0x9b, 1, 0}},
		{FramePacket{Cmd: FrameOpen, Status: FrameError, Channel: 13},
			[]byte{0, 0, 0, 13, 0, 1}},
		{FramePacket{Cmd: FrameData, Channel: 11, Data: []byte("hi")},
			[]byte{0, 2, 0, 11, 2, 0, 'h', 'i'}},
	}

	for _, test := range tests {
		got := test.pkt.Bytes()
		if !reflect.DeepEqual(got, test.exp) {
			t.Errorf("Error encoding %v\nExpected:\n%#v\nGot:\n%#v",
				test.pkt, test.exp, got)
		}
		t.Logf("Packet: %v", test.pkt)
	}
}

func TestErrorStringing(t *testing.T) {
	e := frameError{Status: FrameError, Data: []byte("broken")}
	got := e.Error()
	want := `status=Error, data=broken`
	if got != want {
		t.Errorf("Wanted %v, got %v", want, got)
	}

	e = frameError{Status: 11, Data: []byte("broken and unknown")}
	got = e.Error()
	want = `status={FrameStatus 0xb}, data=broken and unknown`
	if got != want {
		t.Errorf("Wanted %v, got %v", want, got)
	}
}

func benchEncoding(b *testing.B, size int) {
	pkt := FramePacket{
		Cmd:     FrameData,
		Channel: 8184,
		Data:    make([]byte, size),
	}

	b.SetBytes(int64(len(pkt.Bytes())))

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		pkt.Bytes()
	}
}

func BenchmarkEncoding0(b *testing.B) {
	benchEncoding(b, 0)
}

func BenchmarkEncoding8(b *testing.B) {
	benchEncoding(b, 8)
}

func BenchmarkEncoding16(b *testing.B) {
	benchEncoding(b, 16)
}

func BenchmarkEncoding64(b *testing.B) {
	benchEncoding(b, 64)
}

func BenchmarkEncoding256(b *testing.B) {
	benchEncoding(b, 256)
}

func BenchmarkEncoding1024(b *testing.B) {
	benchEncoding(b, 1024)
}

func BenchmarkEncoding8192(b *testing.B) {
	benchEncoding(b, 8192)
}
<|endoftext|>


<commit_before>package handlers

import (
	"errors"
	"io"
	"net/http"

	ctxu "github.com/docker/distribution/context"
	"github.com/docker/distribution/registry/api/errcode"
)

// closeResources closes all the provided resources after running the target
// handler.
func closeResources(handler http.Handler, closers ...io.Closer) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		for _, closer := range closers {
			defer closer.Close()
		}
		handler.ServeHTTP(w, r)
	})
}

// copyFullPayload copies the payload of a HTTP request to destWriter. If it
// receives less content than expected, and the client disconnected during the
// upload, it avoids sending a 400 error to keep the logs cleaner.
func copyFullPayload(responseWriter http.ResponseWriter, r *http.Request, destWriter io.Writer, context ctxu.Context, action string, errSlice *errcode.Errors) error {
	// Get a channel that tells us if the client disconnects
	var clientClosed <-chan bool
	if notifier, ok := responseWriter.(http.CloseNotifier); ok {
		clientClosed = notifier.CloseNotify()
	} else {
		ctxu.GetLogger(context).Warn("the ResponseWriter does not implement CloseNotifier")
	}

	// Read in the data, if any.
	copied, err := io.Copy(destWriter, r.Body)
	if clientClosed != nil && (err != nil || (r.ContentLength > 0 && copied < r.ContentLength)) {
		// Didn't recieve as much content as expected. Did the client
		// disconnect during the request? If so, avoid returning a 400
		// error to keep the logs cleaner.
		select {
		case <-clientClosed:
			// Set the response code to "499 Client Closed Request"
			// Even though the connection has already been closed,
			// this causes the logger to pick up a 499 error
			// instead of showing 0 for the HTTP status.
			responseWriter.WriteHeader(499)

			ctxu.GetLogger(context).Error("client disconnected during " + action)
			return errors.New("client disconnected")
		default:
		}
	}

	if err != nil {
		ctxu.GetLogger(context).Errorf("unknown error reading request payload: %v", err)
		*errSlice = append(*errSlice, errcode.ErrorCodeUnknown.WithDetail(err))
		return err
	}

	return nil
}
<commit_msg>Fix CloseNotifier handling and avoid "the ResponseWriter does not implement CloseNotifier" warnings in logs<commit_after>package handlers

import (
	"errors"
	"io"
	"net/http"

	ctxu "github.com/docker/distribution/context"
	"github.com/docker/distribution/registry/api/errcode"
)

// closeResources closes all the provided resources after running the target
// handler.
func closeResources(handler http.Handler, closers ...io.Closer) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		for _, closer := range closers {
			defer closer.Close()
		}
		handler.ServeHTTP(w, r)
	})
}

// copyFullPayload copies the payload of a HTTP request to destWriter. If it
// receives less content than expected, and the client disconnected during the
// upload, it avoids sending a 400 error to keep the logs cleaner.
func copyFullPayload(responseWriter http.ResponseWriter, r *http.Request, destWriter io.Writer, context ctxu.Context, action string, errSlice *errcode.Errors) error {
	// Get a channel that tells us if the client disconnects
	var clientClosed <-chan bool
	if notifier, ok := responseWriter.(http.CloseNotifier); ok {
		clientClosed = notifier.CloseNotify()
	} else {
		ctxu.GetLogger(context).Warnf("the ResponseWriter does not implement CloseNotifier (type: %T)", responseWriter)
	}

	// Read in the data, if any.
	copied, err := io.Copy(destWriter, r.Body)
	if clientClosed != nil && (err != nil || (r.ContentLength > 0 && copied < r.ContentLength)) {
		// Didn't recieve as much content as expected. Did the client
		// disconnect during the request? If so, avoid returning a 400
		// error to keep the logs cleaner.
		select {
		case <-clientClosed:
			// Set the response code to "499 Client Closed Request"
			// Even though the connection has already been closed,
			// this causes the logger to pick up a 499 error
			// instead of showing 0 for the HTTP status.
			responseWriter.WriteHeader(499)

			ctxu.GetLogger(context).Error("client disconnected during " + action)
			return errors.New("client disconnected")
		default:
		}
	}

	if err != nil {
		ctxu.GetLogger(context).Errorf("unknown error reading request payload: %v", err)
		*errSlice = append(*errSlice, errcode.ErrorCodeUnknown.WithDetail(err))
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package util

import (
	"regexp"
	"strconv"
	"strings"
)

var (
	// stdErrorRegex recognizes the error/warning lines from pyang/confd
	// It currently recognizes the following two patterns:
	// - path:line#:status:message
	// - path:line#(subpath:line#):status:message
	//     NOTE: The subpath info in brackets is currently lumped into one group.
	stdErrorRegex = regexp.MustCompile(`^([^:]+):\s*(\d+)\s*(\([^\)]+\))?\s*:([^:]+):(.+)$`)
)

// StandardErrorLine contains a parsed commandline output from pyang.
type StandardErrorLine struct {
	Path    string
	LineNo  int32
	Status  string
	Message string
}

// StandardOutput contains the parsed commandline outputs from pyang.
type StandardOutput struct {
	ErrorLines   []*StandardErrorLine
	WarningLines []*StandardErrorLine
	OtherLines   []string
}

// ParseStandardOutput parses raw pyang/confd output into a structured format.
// It recognizes two formats of output from pyang and confD:
// <file path>:<line no>:<error/warning>:<message>
// <file path>:<line#>(<import file path>:<line#>):<error/warning>:<message>
func ParseStandardOutput(rawOut string) StandardOutput {
	var out StandardOutput
	for _, line := range strings.Split(rawOut, "\n") {
		if line = strings.TrimSpace(line); line == "" {
			continue
		}

		matches := stdErrorRegex.FindStringSubmatch(line)
		if matches == nil {
			out.OtherLines = append(out.OtherLines, line)
			continue
		}

		filePath := strings.TrimSpace(matches[1])
		lineNumber, err := strconv.ParseInt(strings.TrimSpace(matches[2]), 10, 32)
		if err != nil {
			out.OtherLines = append(out.OtherLines, line)
			continue
		}
		status := strings.ToLower(strings.TrimSpace(matches[4]))
		message := strings.TrimSpace(matches[5])

		switch {
		case strings.Contains(status, "error"):
			out.ErrorLines = append(out.ErrorLines, &StandardErrorLine{
				Path:    filePath,
				LineNo:  int32(lineNumber),
				Status:  status,
				Message: message,
			})
		case strings.Contains(status, "warning"):
			out.WarningLines = append(out.WarningLines, &StandardErrorLine{
				Path:    filePath,
				LineNo:  int32(lineNumber),
				Status:  status,
				Message: message,
			})
		default: // Unrecognized line, so classify as "other".
			out.OtherLines = append(out.OtherLines, line)
		}
	}
	return out
}
<commit_msg>Add package comment<commit_after>// Package util contain utility functions for doing YANG model validation.
package util

import (
	"regexp"
	"strconv"
	"strings"
)

var (
	// stdErrorRegex recognizes the error/warning lines from pyang/confd
	// It currently recognizes the following two patterns:
	// - path:line#:status:message
	// - path:line#(subpath:line#):status:message
	//     NOTE: The subpath info in brackets is currently lumped into one group.
	stdErrorRegex = regexp.MustCompile(`^([^:]+):\s*(\d+)\s*(\([^\)]+\))?\s*:([^:]+):(.+)$`)
)

// StandardErrorLine contains a parsed commandline output from pyang.
type StandardErrorLine struct {
	Path    string
	LineNo  int32
	Status  string
	Message string
}

// StandardOutput contains the parsed commandline outputs from pyang.
type StandardOutput struct {
	ErrorLines   []*StandardErrorLine
	WarningLines []*StandardErrorLine
	OtherLines   []string
}

// ParseStandardOutput parses raw pyang/confd output into a structured format.
// It recognizes two formats of output from pyang and confD:
// <file path>:<line no>:<error/warning>:<message>
// <file path>:<line#>(<import file path>:<line#>):<error/warning>:<message>
func ParseStandardOutput(rawOut string) StandardOutput {
	var out StandardOutput
	for _, line := range strings.Split(rawOut, "\n") {
		if line = strings.TrimSpace(line); line == "" {
			continue
		}

		matches := stdErrorRegex.FindStringSubmatch(line)
		if matches == nil {
			out.OtherLines = append(out.OtherLines, line)
			continue
		}

		filePath := strings.TrimSpace(matches[1])
		lineNumber, err := strconv.ParseInt(strings.TrimSpace(matches[2]), 10, 32)
		if err != nil {
			out.OtherLines = append(out.OtherLines, line)
			continue
		}
		status := strings.ToLower(strings.TrimSpace(matches[4]))
		message := strings.TrimSpace(matches[5])

		switch {
		case strings.Contains(status, "error"):
			out.ErrorLines = append(out.ErrorLines, &StandardErrorLine{
				Path:    filePath,
				LineNo:  int32(lineNumber),
				Status:  status,
				Message: message,
			})
		case strings.Contains(status, "warning"):
			out.WarningLines = append(out.WarningLines, &StandardErrorLine{
				Path:    filePath,
				LineNo:  int32(lineNumber),
				Status:  status,
				Message: message,
			})
		default: // Unrecognized line, so classify as "other".
			out.OtherLines = append(out.OtherLines, line)
		}
	}
	return out
}
<|endoftext|>


<commit_before>package util

import (
	"testing"

	"github.com/golang/protobuf/proto"
	descriptor "github.com/golang/protobuf/protoc-gen-go/descriptor"
	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
)

func TestParseParams(t *testing.T) {
	params := ParseParams(&plugin.CodeGeneratorRequest{
		Parameter: proto.String("key =value,abc = d ef , z = g "),
	})
	if len(params) != 3 {
		t.Fatal("expected 3 arguments got", len(params))
	}
	if params["key"] != "value" {
		t.Fatal(`"key" != "value"`)
	}
	if params["abc"] != "d ef" {
		t.Fatal(`"abc" != "d ef"`)
	}
	if params["z"] != "g" {
		t.Fatal(`"z" != "g"`)
	}
}

func TestIsFullyQualified(t *testing.T) {
	tests := map[string]bool{
		".google.protobuf.UninterpretedOption": true,
		".google.protobuf.FieldOptions.CType":  true,
		"protobuf.FieldOptions.CType":          false,
		"UninterpretedOption":                  false,
	}
	for symbolPath, want := range tests {
		got := IsFullyQualified(symbolPath)
		if got != want {
			t.Fatalf("got %v want %v", got, want)
		}
	}
}

func TestTrimElem(t *testing.T) {
	tests := []struct {
		symbolPath, want string
		n                int
	}{
		// Standard cases:
		{"a.b.c", "b.c", 1},
		{".a.b.c", "b.c", 1},
		{".a.b.c", ".a.b", -1},

		// Extreme cases:
		{"a.b.c", "", 1000},
		{".a.b.c", "", 1000},
		{"a.b.c", "", -1000},
		{".a.b.c", "", -1000},

		{".a.b.c.d", "b.c.d", 1},
		{"a.b.c.d", "b.c.d", 1},
		{".a.b.c.d", "c.d", 2},
		{"a.b.c.d.e", "", 1000},
	}
	for _, tst := range tests {
		got := TrimElem(tst.symbolPath, tst.n)
		if got != tst.want {
			t.Logf("symbolPath=%q\n", tst.symbolPath)
			t.Logf("n=%v\n", tst.n)
			t.Fatalf("got %q want %q\n", got, tst.want)
		}
	}
}

func TestCountElem(t *testing.T) {
	tests := []struct {
		symbolPath string
		want       int
	}{
		{"a.b.c", 3},
		{".a.b.c", 3},
		{"a.b.c.d", 4},
		{"a", 1},
		{".", 0},
		{"", 0},
	}
	for _, tst := range tests {
		got := CountElem(tst.symbolPath)
		if got != tst.want {
			t.Logf("symbolPath=%q\n", tst.symbolPath)
			t.Fatalf("got %v want %v\n", got, tst.want)
		}
	}
}

func TestPackageName(t *testing.T) {
	got := PackageName(&descriptor.FileDescriptorProto{
		Package: proto.String("foo"),
	})
	if got != "foo" {
		t.Fatal("expected explicit package name \"foo\", got %q\n", got)
	}

	got = PackageName(&descriptor.FileDescriptorProto{
		Name: proto.String("some/arbitrary/file.proto"),
	})
	if got != "file" {
		t.Fatalf("expected derived package name \"file\", got %q\n", got)
	}
}
<commit_msg>util: fix Fatal -> Fatalf issue identified with go vet<commit_after>package util

import (
	"testing"

	"github.com/golang/protobuf/proto"
	descriptor "github.com/golang/protobuf/protoc-gen-go/descriptor"
	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
)

func TestParseParams(t *testing.T) {
	params := ParseParams(&plugin.CodeGeneratorRequest{
		Parameter: proto.String("key =value,abc = d ef , z = g "),
	})
	if len(params) != 3 {
		t.Fatal("expected 3 arguments got", len(params))
	}
	if params["key"] != "value" {
		t.Fatal(`"key" != "value"`)
	}
	if params["abc"] != "d ef" {
		t.Fatal(`"abc" != "d ef"`)
	}
	if params["z"] != "g" {
		t.Fatal(`"z" != "g"`)
	}
}

func TestIsFullyQualified(t *testing.T) {
	tests := map[string]bool{
		".google.protobuf.UninterpretedOption": true,
		".google.protobuf.FieldOptions.CType":  true,
		"protobuf.FieldOptions.CType":          false,
		"UninterpretedOption":                  false,
	}
	for symbolPath, want := range tests {
		got := IsFullyQualified(symbolPath)
		if got != want {
			t.Fatalf("got %v want %v", got, want)
		}
	}
}

func TestTrimElem(t *testing.T) {
	tests := []struct {
		symbolPath, want string
		n                int
	}{
		// Standard cases:
		{"a.b.c", "b.c", 1},
		{".a.b.c", "b.c", 1},
		{".a.b.c", ".a.b", -1},

		// Extreme cases:
		{"a.b.c", "", 1000},
		{".a.b.c", "", 1000},
		{"a.b.c", "", -1000},
		{".a.b.c", "", -1000},

		{".a.b.c.d", "b.c.d", 1},
		{"a.b.c.d", "b.c.d", 1},
		{".a.b.c.d", "c.d", 2},
		{"a.b.c.d.e", "", 1000},
	}
	for _, tst := range tests {
		got := TrimElem(tst.symbolPath, tst.n)
		if got != tst.want {
			t.Logf("symbolPath=%q\n", tst.symbolPath)
			t.Logf("n=%v\n", tst.n)
			t.Fatalf("got %q want %q\n", got, tst.want)
		}
	}
}

func TestCountElem(t *testing.T) {
	tests := []struct {
		symbolPath string
		want       int
	}{
		{"a.b.c", 3},
		{".a.b.c", 3},
		{"a.b.c.d", 4},
		{"a", 1},
		{".", 0},
		{"", 0},
	}
	for _, tst := range tests {
		got := CountElem(tst.symbolPath)
		if got != tst.want {
			t.Logf("symbolPath=%q\n", tst.symbolPath)
			t.Fatalf("got %v want %v\n", got, tst.want)
		}
	}
}

func TestPackageName(t *testing.T) {
	got := PackageName(&descriptor.FileDescriptorProto{
		Package: proto.String("foo"),
	})
	if got != "foo" {
		t.Fatalf("expected explicit package name \"foo\", got %q\n", got)
	}

	got = PackageName(&descriptor.FileDescriptorProto{
		Name: proto.String("some/arbitrary/file.proto"),
	})
	if got != "file" {
		t.Fatalf("expected derived package name \"file\", got %q\n", got)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013, Cong Ding. 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.
//
// Author: Cong Ding <dinggnu@gmail.com>

package stun

import (
	"errors"
	"net"
)

// Padding the length of the byte slice to multiple of 4.
func padding(bytes []byte) []byte {
	length := uint16(len(bytes))
	return append(bytes, make([]byte, align(length)-length)...)
}

// Align the uint16 number to the smallest multiple of 4, which is larger than
// or equal to the uint16 number.
func align(n uint16) uint16 {
	return (n + 3) & 0xfffc
}

func sendBindingReq(serverAddr string) (*packet, string, error) {
	conn, err := net.Dial("udp", serverAddr)
	if err != nil {
		return nil, "", err
	}
	// Construct packet.
	packet := newPacket()
	packet.types = type_BINDING_REQUEST
	attribute := newSoftwareAttribute(packet, DefaultSoftwareName)
	packet.addAttribute(*attribute)
	attribute = newFingerprintAttribute(packet)
	packet.addAttribute(*attribute)
	// Send packet.
	localAddr := conn.LocalAddr().String()
	packet, err = packet.send(conn)
	if err != nil {
		return nil, "", err
	}
	err = conn.Close()
	return packet, localAddr, err
}

func sendChangeReq(serverAddr string, changeIP bool, changePort bool) (*packet, error) {
	conn, err := net.Dial("udp", serverAddr)
	if err != nil {
		return nil, err
	}
	// Construct packet.
	packet := newPacket()
	packet.types = type_BINDING_REQUEST
	attribute := newSoftwareAttribute(packet, DefaultSoftwareName)
	packet.addAttribute(*attribute)
	attribute = newChangeReqAttribute(packet, changeIP, changePort)
	packet.addAttribute(*attribute)
	attribute = newFingerprintAttribute(packet)
	packet.addAttribute(*attribute)
	// Send packet.
	packet, err = packet.send(conn)
	if err != nil {
		return nil, err
	}
	err = conn.Close()
	return packet, err
}

func test1(serverAddr string) (*packet, string, bool, *Host, error) {
	packet, localAddr, err := sendBindingReq(serverAddr)
	if err != nil {
		return nil, "", false, nil, err
	}
	if packet == nil {
		return nil, "", false, nil, nil
	}

	// RFC 3489 doesn't require the server return XOR mapped address.
	hostMappedAddr := packet.xorMappedAddr()
	if hostMappedAddr == nil {
		hostMappedAddr = packet.mappedAddr()
		if hostMappedAddr == nil {
			return nil, "", false, nil, errors.New("No mapped address.")
		}
	}

	hostChangedAddr := packet.changedAddr()
	if hostChangedAddr == nil {
		return nil, "", false, nil, errors.New("No changed address.")
	}
	changeAddr := hostChangedAddr.TransportAddr()
	identical := localAddr == hostMappedAddr.TransportAddr()
	return packet, changeAddr, identical, hostMappedAddr, nil
}

func test2(serverAddr string) (*packet, error) {
	return sendChangeReq(serverAddr, true, true)
}

func test3(serverAddr string) (*packet, error) {
	return sendChangeReq(serverAddr, false, true)
}

// Follow RFC 3489 and RFC 5389.
func discover(serverAddr string) (NATType, *Host, error) {
	packet, changeAddr, identical, host, err := test1(serverAddr)
	if err != nil {
		return NAT_ERROR, nil, err
	}
	if packet == nil {
		return NAT_BLOCKED, nil, err
	}
	if identical {
		packet, err = test2(serverAddr)
		if err != nil {
			return NAT_ERROR, host, err
		}
		if packet != nil {
			return NAT_NONE, host, nil
		}
		return NAT_SYMETRIC_UDP_FIREWALL, host, nil
	}
	packet, err = test2(serverAddr)
	if err != nil {
		return NAT_ERROR, host, err
	}
	if packet != nil {
		return NAT_FULL, host, nil
	}
	packet, _, identical, _, err = test1(changeAddr)
	if err != nil {
		return NAT_ERROR, host, err
	}
	if packet == nil {
		// It should be NAT_BLOCKED, but will be detected in the first
		// step. So this will never happen.
		return NAT_UNKNOWN, host, nil
	}
	if identical {
		packet, err = test3(serverAddr)
		if err != nil {
			return NAT_ERROR, host, err
		}
		if packet == nil {
			return NAT_PORT_RESTRICTED, host, nil
		}
		return NAT_RESTRICTED, host, nil
	}
	return NAT_SYMETRIC, host, nil
}
<commit_msg>return nil when err==nil<commit_after>// Copyright 2013, Cong Ding. 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.
//
// Author: Cong Ding <dinggnu@gmail.com>

package stun

import (
	"errors"
	"net"
)

// Padding the length of the byte slice to multiple of 4.
func padding(bytes []byte) []byte {
	length := uint16(len(bytes))
	return append(bytes, make([]byte, align(length)-length)...)
}

// Align the uint16 number to the smallest multiple of 4, which is larger than
// or equal to the uint16 number.
func align(n uint16) uint16 {
	return (n + 3) & 0xfffc
}

func sendBindingReq(serverAddr string) (*packet, string, error) {
	conn, err := net.Dial("udp", serverAddr)
	if err != nil {
		return nil, "", err
	}
	// Construct packet.
	packet := newPacket()
	packet.types = type_BINDING_REQUEST
	attribute := newSoftwareAttribute(packet, DefaultSoftwareName)
	packet.addAttribute(*attribute)
	attribute = newFingerprintAttribute(packet)
	packet.addAttribute(*attribute)
	// Send packet.
	localAddr := conn.LocalAddr().String()
	packet, err = packet.send(conn)
	if err != nil {
		return nil, "", err
	}
	err = conn.Close()
	return packet, localAddr, err
}

func sendChangeReq(serverAddr string, changeIP bool, changePort bool) (*packet, error) {
	conn, err := net.Dial("udp", serverAddr)
	if err != nil {
		return nil, err
	}
	// Construct packet.
	packet := newPacket()
	packet.types = type_BINDING_REQUEST
	attribute := newSoftwareAttribute(packet, DefaultSoftwareName)
	packet.addAttribute(*attribute)
	attribute = newChangeReqAttribute(packet, changeIP, changePort)
	packet.addAttribute(*attribute)
	attribute = newFingerprintAttribute(packet)
	packet.addAttribute(*attribute)
	// Send packet.
	packet, err = packet.send(conn)
	if err != nil {
		return nil, err
	}
	err = conn.Close()
	return packet, err
}

func test1(serverAddr string) (*packet, string, bool, *Host, error) {
	packet, localAddr, err := sendBindingReq(serverAddr)
	if err != nil {
		return nil, "", false, nil, err
	}
	if packet == nil {
		return nil, "", false, nil, nil
	}

	// RFC 3489 doesn't require the server return XOR mapped address.
	hostMappedAddr := packet.xorMappedAddr()
	if hostMappedAddr == nil {
		hostMappedAddr = packet.mappedAddr()
		if hostMappedAddr == nil {
			return nil, "", false, nil, errors.New("No mapped address.")
		}
	}

	hostChangedAddr := packet.changedAddr()
	if hostChangedAddr == nil {
		return nil, "", false, nil, errors.New("No changed address.")
	}
	changeAddr := hostChangedAddr.TransportAddr()
	identical := localAddr == hostMappedAddr.TransportAddr()
	return packet, changeAddr, identical, hostMappedAddr, nil
}

func test2(serverAddr string) (*packet, error) {
	return sendChangeReq(serverAddr, true, true)
}

func test3(serverAddr string) (*packet, error) {
	return sendChangeReq(serverAddr, false, true)
}

// Follow RFC 3489 and RFC 5389.
func discover(serverAddr string) (NATType, *Host, error) {
	packet, changeAddr, identical, host, err := test1(serverAddr)
	if err != nil {
		return NAT_ERROR, nil, err
	}
	if packet == nil {
		return NAT_BLOCKED, nil, nil
	}
	if identical {
		packet, err = test2(serverAddr)
		if err != nil {
			return NAT_ERROR, host, err
		}
		if packet != nil {
			return NAT_NONE, host, nil
		}
		return NAT_SYMETRIC_UDP_FIREWALL, host, nil
	}
	packet, err = test2(serverAddr)
	if err != nil {
		return NAT_ERROR, host, err
	}
	if packet != nil {
		return NAT_FULL, host, nil
	}
	packet, _, identical, _, err = test1(changeAddr)
	if err != nil {
		return NAT_ERROR, host, err
	}
	if packet == nil {
		// It should be NAT_BLOCKED, but will be detected in the first
		// step. So this will never happen.
		return NAT_UNKNOWN, host, nil
	}
	if identical {
		packet, err = test3(serverAddr)
		if err != nil {
			return NAT_ERROR, host, err
		}
		if packet == nil {
			return NAT_PORT_RESTRICTED, host, nil
		}
		return NAT_RESTRICTED, host, nil
	}
	return NAT_SYMETRIC, host, nil
}
<|endoftext|>


<commit_before>// Package server contains logic to handle DNS requests.
package server

import (
	"fmt"
	"strings"

	"github.com/miekg/dns"

	"github.com/nanopack/shaman/config"
	"github.com/nanopack/shaman/core"
)

// Start starts the DNS listener
func Start() error {
	dns.HandleFunc(".", handlerFunc)
	udpListener := &dns.Server{Addr: config.DnsListen, Net: "udp"}
	config.Log.Info("DNS listening at udp://%v", config.DnsListen)
	return fmt.Errorf("DNS listener stopped - %v", udpListener.ListenAndServe())
}

// handlerFunc receives requests, looks up the result and returns what is found.
func handlerFunc(res dns.ResponseWriter, req *dns.Msg) {
	message := new(dns.Msg)
	switch req.Opcode {
	case dns.OpcodeQuery:
		message.SetReply(req)
		message.Compress = false
		message.Answer = make([]dns.RR, 0)

		for _, question := range message.Question {
			answers := answerQuestion(question)
			for i := range answers {
				message.Answer = append(message.Answer, answers[i])
			}
		}
		if len(message.Answer) == 0 {
			message.Rcode = dns.RcodeNameError
		}
	default:
		message = message.SetRcode(req, dns.RcodeNotImplemented)
	}
	res.WriteMsg(message)
}

// answerQuestion returns resource record answers for the domain in question
func answerQuestion(question dns.Question) []dns.RR {
	answers := make([]dns.RR, 0)

	// get the resource (check memory, cache, and (todo:) upstream)
	r, err := shaman.GetRecord(question.Name)
	if err != nil {
		config.Log.Trace("Failed to get records for '%s' - %v", question.Name, err)
	}

	// validate the records and append correct type to answers[]
	for _, record := range r.StringSlice() {
		entry, err := dns.NewRR(record)
		if err != nil {
			config.Log.Debug("Failed to create RR from record - %v", err)
			continue
		}
		entry.Header().Name = question.Name
		if entry.Header().Rrtype == question.Qtype || question.Qtype == dns.TypeANY {
			answers = append(answers, entry)
		}
	}

	// todo: should `shaman.GetRecord` be wildcard aware (*.domain.com) or is this ok
	// recursively resolve if no records found
	if len(answers) == 0 {
		question.Name = stripSubdomain(question.Name)
		if len(question.Name) > 0 {
			config.Log.Trace("Checking again with '%v'", question.Name)
			return answerQuestion(question)
		}
	}

	return answers
}

// stripSubdomain strips off the subbest domain, returning the domain (won't return TLD)
func stripSubdomain(name string) string {
	words := 3 // assume rooted domain (end with '.')
	// handle edge case of unrooted domain
	t := []byte(name)
	if len(t) > 0 && t[len(t)-1] != '.' {
		words = 2
	}

	config.Log.Trace("Stripping subdomain from '%v'", name)
	names := strings.Split(name, ".")

	// prevent searching for just 'com.' (["domain", "com", ""])
	if len(names) > words {
		return strings.Join(names[1:], ".")
	}
	return ""
}
<commit_msg>if it cant find a record, return the SOA as part of the Ns on the message. If there is no answers and no SOA, then return the NXDOMAIN error.<commit_after>// Package server contains logic to handle DNS requests.
package server

import (
	"fmt"
	"strings"

	"github.com/miekg/dns"

	"github.com/nanopack/shaman/config"
	"github.com/nanopack/shaman/core"
)

// Start starts the DNS listener
func Start() error {
	dns.HandleFunc(".", handlerFunc)
	udpListener := &dns.Server{Addr: config.DnsListen, Net: "udp"}
	config.Log.Info("DNS listening at udp://%v", config.DnsListen)
	return fmt.Errorf("DNS listener stopped - %v", udpListener.ListenAndServe())
}

// handlerFunc receives requests, looks up the result and returns what is found.
func handlerFunc(res dns.ResponseWriter, req *dns.Msg) {
	message := new(dns.Msg)
	switch req.Opcode {
	case dns.OpcodeQuery:
		message.SetReply(req)
		message.Compress = false
		message.Answer = make([]dns.RR, 0)

		for _, question := range message.Question {
			answers := answerQuestion(question.Name, question.Qtype)
			for i := range answers {
				message.Answer = append(message.Answer, answers[i])
			}
		}
		if (len(message.Answer) == 0 ){
			// If there are no records, go back through and search for SOA records
			for _, question := range message.Question {
				answers := answerQuestion(question.Name, dns.TypeSOA)
				for i := range answers {
					message.Ns = append(message.Ns, answers[i])
				}
			}
		}

		if (len(message.Answer) == 0  && len(message.Ns) == 0 ){
			message.Rcode = dns.RcodeNameError
		}
	default:
		message = message.SetRcode(req, dns.RcodeNotImplemented)
	}
	res.WriteMsg(message)
}

// answerQuestion returns resource record answers for the domain in question
func answerQuestion(name string, qtype uint16) []dns.RR {
	answers := make([]dns.RR, 0)

	// get the resource (check memory, cache, and (todo:) upstream)
	r, err := shaman.GetRecord(name)
	if err != nil {
		config.Log.Trace("Failed to get records for '%s' - %v", name, err)
	}

	// validate the records and append correct type to answers[]
	for _, record := range r.StringSlice() {
		entry, err := dns.NewRR(record)
		if err != nil {
			config.Log.Debug("Failed to create RR from record - %v", err)
			continue
		}
		entry.Header().Name = name
		if entry.Header().Rrtype == qtype || qtype == dns.TypeANY {
			answers = append(answers, entry)
		}
	}

	// todo: should `shaman.GetRecord` be wildcard aware (*.domain.com) or is this ok
	// recursively resolve if no records found
	if len(answers) == 0 {
		name = stripSubdomain(name)
		if len(name) > 0 {
			config.Log.Trace("Checking again with '%v'", name)
			return answerQuestion(name, qtype)
		}
	}

	return answers
}

// stripSubdomain strips off the subbest domain, returning the domain (won't return TLD)
func stripSubdomain(name string) string {
	words := 3 // assume rooted domain (end with '.')
	// handle edge case of unrooted domain
	t := []byte(name)
	if len(t) > 0 && t[len(t)-1] != '.' {
		words = 2
	}

	config.Log.Trace("Stripping subdomain from '%v'", name)
	names := strings.Split(name, ".")

	// prevent searching for just 'com.' (["domain", "com", ""])
	if len(names) > words {
		return strings.Join(names[1:], ".")
	}
	return ""
}
<|endoftext|>


<commit_before>package brain

import (
	"io"

	"github.com/BytemarkHosting/bytemark-client/lib/output"
	"github.com/BytemarkHosting/bytemark-client/lib/output/prettyprint"
)

// MigrationJobs represents more than one account in output.Outputtable form.
type MigrationJobs []MigrationJob

// DefaultFields returns the list of default fields to feed to github.com/BytemarkHosting/row.From for this type, which is the same as MigrationJob.DefaultFields.
func (mjs MigrationJobs) DefaultFields(f output.Format) string {
	return (MigrationJob{}).DefaultFields(f)
}

// PrettyPrint writes a human-readable summary of the migration jobs to writer at the given detail level.
func (mjs MigrationJobs) PrettyPrint(wr io.Writer, detail prettyprint.DetailLevel) error {
	migrationjobsTpl := `
{{ define "migrationjobs_sgl" }}{{ len . }} servers{{ end }}

{{ define "migrationjobs_medium" -}}
{{- range . -}}
{{- prettysprint . "_sgl" }}
{{ end -}}
{{- end }}

{{ define "migrationjobs_full" }}{{ template "migrationjobs_medium" . }}{{ end }}
`
	return prettyprint.Run(wr, migrationjobsTpl, "migrationjobs"+string(detail), mjs)
}
<commit_msg>fix comment<commit_after>package brain

import (
	"io"

	"github.com/BytemarkHosting/bytemark-client/lib/output"
	"github.com/BytemarkHosting/bytemark-client/lib/output/prettyprint"
)

// MigrationJobs represents more than one migration job in output.Outputtable form.
type MigrationJobs []MigrationJob

// DefaultFields returns the list of default fields to feed to github.com/BytemarkHosting/row.From for this type, which is the same as MigrationJob.DefaultFields.
func (mjs MigrationJobs) DefaultFields(f output.Format) string {
	return (MigrationJob{}).DefaultFields(f)
}

// PrettyPrint writes a human-readable summary of the migration jobs to writer at the given detail level.
func (mjs MigrationJobs) PrettyPrint(wr io.Writer, detail prettyprint.DetailLevel) error {
	migrationjobsTpl := `
{{ define "migrationjobs_sgl" }}{{ len . }} servers{{ end }}

{{ define "migrationjobs_medium" -}}
{{- range . -}}
{{- prettysprint . "_sgl" }}
{{ end -}}
{{- end }}

{{ define "migrationjobs_full" }}{{ template "migrationjobs_medium" . }}{{ end }}
`
	return prettyprint.Run(wr, migrationjobsTpl, "migrationjobs"+string(detail), mjs)
}
<|endoftext|>


<commit_before>package util

import (
	"time"

	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/wait"
	authorizationtypedclient "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/typed/authorization/internalversion"

	"k8s.io/kubernetes/pkg/apis/authorization"
)

const (
	PolicyCachePollInterval = 100 * time.Millisecond
	PolicyCachePollTimeout  = 5 * time.Second
)

// WaitForPolicyUpdate checks if the given client can perform the named verb and action.
// If PolicyCachePollTimeout is reached without the expected condition matching, an error is returned
func WaitForPolicyUpdate(c authorizationtypedclient.SelfSubjectAccessReviewsGetter, namespace, verb string, resource schema.GroupResource, allowed bool) error {
	review := &authorization.SelfSubjectAccessReview{
		Spec: authorization.SelfSubjectAccessReviewSpec{
			ResourceAttributes: &authorization.ResourceAttributes{
				Namespace: namespace,
				Verb:      verb,
				Group:     resource.Group,
				Resource:  resource.Resource,
			},
		},
	}
	err := wait.Poll(PolicyCachePollInterval, PolicyCachePollTimeout, func() (bool, error) {
		response, err := c.SelfSubjectAccessReviews().Create(review)
		if err != nil {
			return false, err
		}
		return response.Status.Allowed == allowed, nil
	})
	return err
}

// WaitForClusterPolicyUpdate checks if the given client can perform the named verb and action.
// If PolicyCachePollTimeout is reached without the expected condition matching, an error is returned
func WaitForClusterPolicyUpdate(c authorizationtypedclient.SelfSubjectAccessReviewsGetter, verb string, resource schema.GroupResource, allowed bool) error {
	review := &authorization.SelfSubjectAccessReview{
		Spec: authorization.SelfSubjectAccessReviewSpec{
			ResourceAttributes: &authorization.ResourceAttributes{
				Verb:     verb,
				Group:    resource.Group,
				Resource: resource.Resource,
			},
		},
	}
	err := wait.Poll(PolicyCachePollInterval, PolicyCachePollTimeout, func() (bool, error) {
		response, err := c.SelfSubjectAccessReviews().Create(review)
		if err != nil {
			return false, err
		}
		if response.Status.Allowed != allowed {
			return false, nil
		}
		return true, nil
	})
	return err
}
<commit_msg>Bump policy timeouts<commit_after>package util

import (
	"time"

	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/wait"
	authorizationtypedclient "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/typed/authorization/internalversion"

	"k8s.io/kubernetes/pkg/apis/authorization"
)

const (
	PolicyCachePollInterval = 100 * time.Millisecond
	PolicyCachePollTimeout  = 10 * time.Second
)

// WaitForPolicyUpdate checks if the given client can perform the named verb and action.
// If PolicyCachePollTimeout is reached without the expected condition matching, an error is returned
func WaitForPolicyUpdate(c authorizationtypedclient.SelfSubjectAccessReviewsGetter, namespace, verb string, resource schema.GroupResource, allowed bool) error {
	review := &authorization.SelfSubjectAccessReview{
		Spec: authorization.SelfSubjectAccessReviewSpec{
			ResourceAttributes: &authorization.ResourceAttributes{
				Namespace: namespace,
				Verb:      verb,
				Group:     resource.Group,
				Resource:  resource.Resource,
			},
		},
	}
	err := wait.Poll(PolicyCachePollInterval, PolicyCachePollTimeout, func() (bool, error) {
		response, err := c.SelfSubjectAccessReviews().Create(review)
		if err != nil {
			return false, err
		}
		return response.Status.Allowed == allowed, nil
	})
	return err
}

// WaitForClusterPolicyUpdate checks if the given client can perform the named verb and action.
// If PolicyCachePollTimeout is reached without the expected condition matching, an error is returned
func WaitForClusterPolicyUpdate(c authorizationtypedclient.SelfSubjectAccessReviewsGetter, verb string, resource schema.GroupResource, allowed bool) error {
	review := &authorization.SelfSubjectAccessReview{
		Spec: authorization.SelfSubjectAccessReviewSpec{
			ResourceAttributes: &authorization.ResourceAttributes{
				Verb:     verb,
				Group:    resource.Group,
				Resource: resource.Resource,
			},
		},
	}
	err := wait.Poll(PolicyCachePollInterval, PolicyCachePollTimeout, func() (bool, error) {
		response, err := c.SelfSubjectAccessReviews().Create(review)
		if err != nil {
			return false, err
		}
		if response.Status.Allowed != allowed {
			return false, nil
		}
		return true, nil
	})
	return err
}
<|endoftext|>


<commit_before>// Copyright 2019 The Wuffs Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://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 cgozstd

import (
	"bytes"
	"io/ioutil"
	"strings"
	"testing"
)

const (
	// compressedMore is 15 bytes of zstd frame:
	// \x28\xb5\x2f\xfd \x00 \x58 \x31\x00\x00 \x4d\x6f\x72\x65\x21\x0a"
	// Magic----------- FHD- WD-- BH---------- BlockData----------------
	//
	// Frame Header Descriptor 0x00 means no flags set.
	// Window Descriptor: Exponent=11, Mantissa=0, Window_Size=2MiB.
	// Block Header: Last_Block=1, Block_Type=0 (Raw_Block), Block_Size=6.
	// BlockData is the literal bytes "More!\n".
	compressedMore = "\x28\xb5\x2f\xfd\x00\x58\x31\x00\x00\x4d\x6f\x72\x65\x21\x0a"

	uncompressedMore = "More!\n"
)

func TestRoundTrip(t *testing.T) {
	if !cgoEnabled {
		t.Skip("cgo is not enabled")
	}

	wr := &WriterRecycler{}
	defer wr.Close()
	w := &Writer{}
	wr.Bind(w)

	rr := &ReaderRecycler{}
	defer rr.Close()
	r := &Reader{}
	rr.Bind(r)

	for i := 0; i < 3; i++ {
		buf := &bytes.Buffer{}

		// Compress.
		{
			w.Reset(buf, nil, 0)
			if _, err := w.Write([]byte(uncompressedMore)); err != nil {
				w.Close()
				t.Fatalf("i=%d: Write: %v", i, err)
			}
			if err := w.Close(); err != nil {
				t.Fatalf("i=%d: Close: %v", i, err)
			}
		}

		compressed := buf.String()
		if compressed != compressedMore {
			t.Fatalf("i=%d: compressed\ngot  % 02x\nwant % 02x", i, compressed, compressedMore)
		}

		// Uncompress.
		{
			r.Reset(strings.NewReader(compressed), nil)
			gotBytes, err := ioutil.ReadAll(r)
			if err != nil {
				r.Close()
				t.Fatalf("i=%d: ReadAll: %v", i, err)
			}
			if got, want := string(gotBytes), uncompressedMore; got != want {
				r.Close()
				t.Fatalf("i=%d:\ngot  %q\nwant %q", i, got, want)
			}
			if err := r.Close(); err != nil {
				t.Fatalf("i=%d: Close: %v", i, err)
			}
		}
	}
}
<commit_msg>Fix cgozstd_test.go comment typo<commit_after>// Copyright 2019 The Wuffs Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://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 cgozstd

import (
	"bytes"
	"io/ioutil"
	"strings"
	"testing"
)

const (
	// compressedMore is 15 bytes of zstd frame:
	// \x28\xb5\x2f\xfd \x00 \x58 \x31\x00\x00 \x4d\x6f\x72\x65\x21\x0a
	// Magic----------- FHD- WD-- BH---------- BlockData---------------
	//
	// Frame Header Descriptor: no flags set.
	// Window Descriptor: Exponent=11, Mantissa=0, Window_Size=2MiB.
	// Block Header: Last_Block=1, Block_Type=0 (Raw_Block), Block_Size=6.
	// Block Data: the literal bytes "More!\n".
	compressedMore = "\x28\xb5\x2f\xfd\x00\x58\x31\x00\x00\x4d\x6f\x72\x65\x21\x0a"

	uncompressedMore = "More!\n"
)

func TestRoundTrip(t *testing.T) {
	if !cgoEnabled {
		t.Skip("cgo is not enabled")
	}

	wr := &WriterRecycler{}
	defer wr.Close()
	w := &Writer{}
	wr.Bind(w)

	rr := &ReaderRecycler{}
	defer rr.Close()
	r := &Reader{}
	rr.Bind(r)

	for i := 0; i < 3; i++ {
		buf := &bytes.Buffer{}

		// Compress.
		{
			w.Reset(buf, nil, 0)
			if _, err := w.Write([]byte(uncompressedMore)); err != nil {
				w.Close()
				t.Fatalf("i=%d: Write: %v", i, err)
			}
			if err := w.Close(); err != nil {
				t.Fatalf("i=%d: Close: %v", i, err)
			}
		}

		compressed := buf.String()
		if compressed != compressedMore {
			t.Fatalf("i=%d: compressed\ngot  % 02x\nwant % 02x", i, compressed, compressedMore)
		}

		// Uncompress.
		{
			r.Reset(strings.NewReader(compressed), nil)
			gotBytes, err := ioutil.ReadAll(r)
			if err != nil {
				r.Close()
				t.Fatalf("i=%d: ReadAll: %v", i, err)
			}
			if got, want := string(gotBytes), uncompressedMore; got != want {
				r.Close()
				t.Fatalf("i=%d:\ngot  %q\nwant %q", i, got, want)
			}
			if err := r.Close(); err != nil {
				t.Fatalf("i=%d: Close: %v", i, err)
			}
		}
	}
}
<|endoftext|>


<commit_before>package cache

import (
	"strconv"
	"testing"
	"time"

	"github.com/mholt/caddy"
	"github.com/stretchr/testify/require"
)

func TestParsingConfig(t *testing.T) {
	tests := []struct {
		input     string
		shouldErr bool
		expect    Config
	}{
		{"cache", false, Config{
			StatusHeader:  defaultStatusHeader,
			LockTimeout:   defaultLockTimeout,
			DefaultMaxAge: defaultMaxAge,
			CacheRules:    []CacheRule{},
		}},
		{"cache {\n match_path /assets \n} }", false, Config{
			StatusHeader:  defaultStatusHeader,
			LockTimeout:   defaultLockTimeout,
			DefaultMaxAge: defaultMaxAge,
			CacheRules:    []CacheRule{&PathCacheRule{Path: "/assets"}},
		}},
		{"cache {\n match_path /assets \n match_path /api \n} \n}", false, Config{
			StatusHeader:  defaultStatusHeader,
			LockTimeout:   defaultLockTimeout,
			DefaultMaxAge: defaultMaxAge,
			CacheRules: []CacheRule{
				&PathCacheRule{Path: "/assets"},
				&PathCacheRule{Path: "/api"},
			},
		}},
		{"cache {\n match_header Content-Type image/png image/gif \n match_path /assets \n}", false, Config{
			StatusHeader:  defaultStatusHeader,
			LockTimeout:   defaultLockTimeout,
			DefaultMaxAge: defaultMaxAge,
			CacheRules: []CacheRule{
				&HeaderCacheRule{Header: "Content-Type", Value: []string{"image/png", "image/gif"}},
				&PathCacheRule{Path: "/assets"},
			},
		}},
		{"cache {\n status_header X-Custom-Header \n}", false, Config{
			StatusHeader:  "X-Custom-Header",
			LockTimeout:   defaultLockTimeout,
			DefaultMaxAge: defaultMaxAge,
			CacheRules:    []CacheRule{},
		}},
		{"cache {\n path /tmp/caddy \n}", false, Config{
			StatusHeader:  defaultStatusHeader,
			LockTimeout:   defaultLockTimeout,
			DefaultMaxAge: defaultMaxAge,
			CacheRules:    []CacheRule{},
			Path:          "/tmp/caddy",
		}},
		{"cache {\n lock_timeout 1s \n}", false, Config{
			StatusHeader:  defaultStatusHeader,
			LockTimeout:   time.Duration(1) * time.Second,
			DefaultMaxAge: defaultMaxAge,
			CacheRules:    []CacheRule{},
		}},
		{"cache {\n default_max_age 1h \n}", false, Config{
			StatusHeader:  defaultStatusHeader,
			LockTimeout:   defaultLockTimeout,
			DefaultMaxAge: time.Duration(1) * time.Hour,
			CacheRules:    []CacheRule{},
		}},
		{"cache {\n match_header aheader \n}", true, Config{}},          // match_header without value
		{"cache {\n lock_timeout aheader \n}", true, Config{}},          // lock_timeout with invalid duration
		{"cache {\n lock_timeout \n}", true, Config{}},                  // lock_timeout has no arguments
		{"cache {\n default_max_age somevalue \n}", true, Config{}},     // lock_timeout has invalid duration
		{"cache {\n default_max_age \n}", true, Config{}},               // default_max_age has no arguments
		{"cache {\n status_header aheader another \n}", true, Config{}}, // status_header with invalid number of parameters
		{"cache {\n match_path / ea \n}", true, Config{}},               // Invalid number of parameters in match
		{"cache {\n invalid / ea \n}", true, Config{}},                  // Invalid directive
		{"cache {\n path \n}", true, Config{}},                          // Path without arguments
	}

	for i, test := range tests {
		t.Run(strconv.Itoa(i), func(t *testing.T) {
			c := caddy.NewTestController("http", test.input)
			actual, err := cacheParse(c)

			if test.shouldErr {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
				require.Equal(t, test.expect, *actual, "Invalid config parsed in test "+strconv.Itoa(i+1))
			}
		})
	}

}
<commit_msg>fix: fixed tests, added some more<commit_after>package cache

import (
	"strconv"
	"testing"
	"time"

	"github.com/mholt/caddy"
	"github.com/stretchr/testify/require"
)

func TestParsingConfig(t *testing.T) {
	tests := []struct {
		input     string
		shouldErr bool
		expect    Config
	}{
		{"cache", false, Config{
			StatusHeader:     defaultStatusHeader,
			LockTimeout:      defaultLockTimeout,
			DefaultMaxAge:    defaultMaxAge,
			CacheRules:       []CacheRule{},
			CacheKeyTemplate: defaultCacheKeyTemplate,
		}},
		{"cache {\n match_path /assets \n} }", false, Config{
			StatusHeader:     defaultStatusHeader,
			LockTimeout:      defaultLockTimeout,
			DefaultMaxAge:    defaultMaxAge,
			CacheRules:       []CacheRule{&PathCacheRule{Path: "/assets"}},
			CacheKeyTemplate: defaultCacheKeyTemplate,
		}},
		{"cache {\n match_path /assets \n match_path /api \n} \n}", false, Config{
			StatusHeader:  defaultStatusHeader,
			LockTimeout:   defaultLockTimeout,
			DefaultMaxAge: defaultMaxAge,
			CacheRules: []CacheRule{
				&PathCacheRule{Path: "/assets"},
				&PathCacheRule{Path: "/api"},
			},
			CacheKeyTemplate: defaultCacheKeyTemplate,
		}},
		{"cache {\n match_header Content-Type image/png image/gif \n match_path /assets \n}", false, Config{
			StatusHeader:  defaultStatusHeader,
			LockTimeout:   defaultLockTimeout,
			DefaultMaxAge: defaultMaxAge,
			CacheRules: []CacheRule{
				&HeaderCacheRule{Header: "Content-Type", Value: []string{"image/png", "image/gif"}},
				&PathCacheRule{Path: "/assets"},
			},
			CacheKeyTemplate: defaultCacheKeyTemplate,
		}},
		{"cache {\n status_header X-Custom-Header \n}", false, Config{
			StatusHeader:     "X-Custom-Header",
			LockTimeout:      defaultLockTimeout,
			DefaultMaxAge:    defaultMaxAge,
			CacheRules:       []CacheRule{},
			CacheKeyTemplate: defaultCacheKeyTemplate,
		}},
		{"cache {\n path /tmp/caddy \n}", false, Config{
			StatusHeader:     defaultStatusHeader,
			LockTimeout:      defaultLockTimeout,
			DefaultMaxAge:    defaultMaxAge,
			CacheRules:       []CacheRule{},
			Path:             "/tmp/caddy",
			CacheKeyTemplate: defaultCacheKeyTemplate,
		}},
		{"cache {\n lock_timeout 1s \n}", false, Config{
			StatusHeader:     defaultStatusHeader,
			LockTimeout:      time.Duration(1) * time.Second,
			DefaultMaxAge:    defaultMaxAge,
			CacheRules:       []CacheRule{},
			CacheKeyTemplate: defaultCacheKeyTemplate,
		}},
		{"cache {\n default_max_age 1h \n}", false, Config{
			StatusHeader:     defaultStatusHeader,
			LockTimeout:      defaultLockTimeout,
			DefaultMaxAge:    time.Duration(1) * time.Hour,
			CacheRules:       []CacheRule{},
			CacheKeyTemplate: defaultCacheKeyTemplate,
		}},
		{"cache {\n cache_key \"{scheme} {host}{uri}\" \n}", false, Config{
			StatusHeader:     defaultStatusHeader,
			LockTimeout:      defaultLockTimeout,
			DefaultMaxAge:    defaultMaxAge,
			CacheRules:       []CacheRule{},
			CacheKeyTemplate: "{scheme} {host}{uri}",
		}},
		{"cache {\n match_header aheader \n}", true, Config{}},          // match_header without value
		{"cache {\n lock_timeout aheader \n}", true, Config{}},          // lock_timeout with invalid duration
		{"cache {\n lock_timeout \n}", true, Config{}},                  // lock_timeout has no arguments
		{"cache {\n default_max_age somevalue \n}", true, Config{}},     // lock_timeout has invalid duration
		{"cache {\n default_max_age \n}", true, Config{}},               // default_max_age has no arguments
		{"cache {\n status_header aheader another \n}", true, Config{}}, // status_header with invalid number of parameters
		{"cache {\n match_path / ea \n}", true, Config{}},               // Invalid number of parameters in match
		{"cache {\n invalid / ea \n}", true, Config{}},                  // Invalid directive
		{"cache {\n path \n}", true, Config{}},                          // Path without arguments
		{"cache {\n cache_key \n}", true, Config{}},                     // cache_key without arguments
	}

	for i, test := range tests {
		t.Run(strconv.Itoa(i), func(t *testing.T) {
			c := caddy.NewTestController("http", test.input)
			actual, err := cacheParse(c)

			if test.shouldErr {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
				require.Equal(t, test.expect, *actual, "Invalid config parsed in test "+strconv.Itoa(i+1))
			}
		})
	}

}
<|endoftext|>


<commit_before><commit_msg>Add base method<commit_after><|endoftext|>


<commit_before>package pms

import (
	"github.com/ambientsound/pms/console"
	"github.com/ambientsound/pms/message"
)

func (pms *PMS) Main() {
	for {
		select {
		case <-pms.QuitSignal:
			pms.handleQuitSignal()
			return
		case <-pms.EventLibrary:
			pms.handleEventLibrary()
		case <-pms.EventQueue:
			pms.handleEventQueue()
		case <-pms.EventIndex:
			pms.handleEventIndex()
		case <-pms.EventList:
			pms.handleEventList()
		case <-pms.EventPlayer:
			pms.handleEventPlayer()
		case key := <-pms.EventOption:
			pms.handleEventOption(key)
		case msg := <-pms.EventMessage:
			pms.handleEventMessage(msg)
		case ev := <-pms.UI.EventKeyInput:
			pms.KeyInput(ev)
		case s := <-pms.UI.EventInputCommand:
			pms.Execute(s)
		}

		// Draw missing parts after every iteration
		pms.UI.App.PostFunc(func() {
			pms.UI.App.Update()
		})
	}
}

func (pms *PMS) handleQuitSignal() {
	console.Log("Received quit signal, exiting.")
	pms.UI.Quit()
}

func (pms *PMS) handleEventLibrary() {
	console.Log("Song library updated in MPD, assigning to UI")
	pms.UI.App.PostFunc(func() {
		pms.UI.Songlist.ReplaceSonglist(pms.Library)
	})
}

func (pms *PMS) handleEventQueue() {
	console.Log("Queue updated in MPD, assigning to UI")
	pms.UI.App.PostFunc(func() {
		pms.UI.Songlist.ReplaceSonglist(pms.Queue)
	})
}

func (pms *PMS) handleEventIndex() {
	console.Log("Search index updated, assigning to UI")
	pms.UI.App.PostFunc(func() {
		pms.UI.SetIndex(pms.Index)
	})
}

func (pms *PMS) handleEventList() {
	console.Log("Songlist changed, notifying UI")
	pms.UI.App.PostFunc(func() {
		pms.UI.Songlist.ListChanged()
	})
}

func (pms *PMS) handleEventOption(key string) {
	console.Log("Option '%s' has been changed", key)
	switch key {
	case "topbar":
		pms.setupTopbar()
	}
}

func (pms *PMS) handleEventPlayer() {
	pms.UI.App.PostFunc(func() {
		pms.UI.Playbar.SetPlayerStatus(pms.CurrentPlayerStatus())
		pms.UI.Playbar.SetSong(pms.CurrentSong())
		pms.UI.Songlist.SetCurrentSong(pms.CurrentSong())
	})
}

func (pms *PMS) handleEventMessage(msg message.Message) {
	message.Log(msg)
	pms.UI.App.PostFunc(func() {
		pms.UI.Multibar.SetMessage(msg)
	})
}
<commit_msg>Set column headers immediately when the 'columns' option is changed<commit_after>package pms

import (
	"github.com/ambientsound/pms/console"
	"github.com/ambientsound/pms/message"
)

func (pms *PMS) Main() {
	for {
		select {
		case <-pms.QuitSignal:
			pms.handleQuitSignal()
			return
		case <-pms.EventLibrary:
			pms.handleEventLibrary()
		case <-pms.EventQueue:
			pms.handleEventQueue()
		case <-pms.EventIndex:
			pms.handleEventIndex()
		case <-pms.EventList:
			pms.handleEventList()
		case <-pms.EventPlayer:
			pms.handleEventPlayer()
		case key := <-pms.EventOption:
			pms.handleEventOption(key)
		case msg := <-pms.EventMessage:
			pms.handleEventMessage(msg)
		case ev := <-pms.UI.EventKeyInput:
			pms.KeyInput(ev)
		case s := <-pms.UI.EventInputCommand:
			pms.Execute(s)
		}

		// Draw missing parts after every iteration
		pms.UI.App.PostFunc(func() {
			pms.UI.App.Update()
		})
	}
}

func (pms *PMS) handleQuitSignal() {
	console.Log("Received quit signal, exiting.")
	pms.UI.Quit()
}

func (pms *PMS) handleEventLibrary() {
	console.Log("Song library updated in MPD, assigning to UI")
	pms.UI.App.PostFunc(func() {
		pms.UI.Songlist.ReplaceSonglist(pms.Library)
	})
}

func (pms *PMS) handleEventQueue() {
	console.Log("Queue updated in MPD, assigning to UI")
	pms.UI.App.PostFunc(func() {
		pms.UI.Songlist.ReplaceSonglist(pms.Queue)
	})
}

func (pms *PMS) handleEventIndex() {
	console.Log("Search index updated, assigning to UI")
	pms.UI.App.PostFunc(func() {
		pms.UI.SetIndex(pms.Index)
	})
}

func (pms *PMS) handleEventList() {
	console.Log("Songlist changed, notifying UI")
	pms.UI.App.PostFunc(func() {
		pms.UI.Songlist.ListChanged()
	})
}

func (pms *PMS) handleEventOption(key string) {
	console.Log("Option '%s' has been changed", key)
	switch key {
	case "topbar":
		pms.setupTopbar()
	case "columns":
		pms.UI.App.PostFunc(func() {
			pms.UI.Songlist.ListChanged()
		})
	}
}

func (pms *PMS) handleEventPlayer() {
	pms.UI.App.PostFunc(func() {
		pms.UI.Playbar.SetPlayerStatus(pms.CurrentPlayerStatus())
		pms.UI.Playbar.SetSong(pms.CurrentSong())
		pms.UI.Songlist.SetCurrentSong(pms.CurrentSong())
	})
}

func (pms *PMS) handleEventMessage(msg message.Message) {
	message.Log(msg)
	pms.UI.App.PostFunc(func() {
		pms.UI.Multibar.SetMessage(msg)
	})
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net"
)

const PORT = "2525"

func main() {
	server := createServer(PORT)

	defer server.Close()

	for {
		handleConnection(acceptConnection(server))
	}
}

func createServer(port string) (net.Listener) {
	server, err := net.Listen("tcp", ":" + PORT)

	if err != nil {
		log.Fatalf("Error listening to port %s\n", PORT)
	}

	return server
}

func acceptConnection(server net.Listener) net.Conn {
	conn, err := server.Accept()

	if err != nil {
		log.Fatalf("Error accepting: %v\n", err)
	}

	return conn
}

func handleConnection(conn net.Conn) {
}
<commit_msg>Remove unnecessary function.<commit_after>package main

import (
	"log"
	"net"
)

const PORT = "2525"

func main() {
	server := createServer(PORT)

	defer server.Close()

	for {
		conn, err := server.Accept()

		if err != nil {
			log.Fatalf("Error accepting: %v\n", err)
		}

		handleConnection(conn)
	}
}

func createServer(port string) (net.Listener) {
	server, err := net.Listen("tcp", ":" + PORT)

	if err != nil {
		log.Fatalf("Error listening to port %s\n", PORT)
	}

	return server
}

func handleConnection(conn net.Conn) {
	defer conn.Close()
}
<|endoftext|>


<commit_before>package aws

import (
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"log"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsIAMServerCertificate() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsIAMServerCertificateCreate,
		Read:   resourceAwsIAMServerCertificateRead,
		Delete: resourceAwsIAMServerCertificateDelete,
		Importer: &schema.ResourceImporter{
			State: resourceAwsIAMServerCertificateImport,
		},

		Schema: map[string]*schema.Schema{
			"certificate_body": {
				Type:      schema.TypeString,
				Required:  true,
				ForceNew:  true,
				StateFunc: normalizeCert,
			},

			"certificate_chain": {
				Type:      schema.TypeString,
				Optional:  true,
				ForceNew:  true,
				StateFunc: normalizeCert,
			},

			"path": {
				Type:     schema.TypeString,
				Optional: true,
				Default:  "/",
				ForceNew: true,
			},

			"private_key": {
				Type:      schema.TypeString,
				Required:  true,
				ForceNew:  true,
				StateFunc: normalizeCert,
				Sensitive: true,
			},

			"name": {
				Type:          schema.TypeString,
				Optional:      true,
				Computed:      true,
				ForceNew:      true,
				ConflictsWith: []string{"name_prefix"},
				ValidateFunc:  validation.StringLenBetween(0, 128),
			},

			"name_prefix": {
				Type:          schema.TypeString,
				Optional:      true,
				ForceNew:      true,
				ConflictsWith: []string{"name"},
				ValidateFunc:  validation.StringLenBetween(0, 128-resource.UniqueIDSuffixLength),
			},

			"arn": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
		},
	}
}

func resourceAwsIAMServerCertificateCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).iamconn

	var sslCertName string
	if v, ok := d.GetOk("name"); ok {
		sslCertName = v.(string)
	} else if v, ok := d.GetOk("name_prefix"); ok {
		sslCertName = resource.PrefixedUniqueId(v.(string))
	} else {
		sslCertName = resource.UniqueId()
	}

	createOpts := &iam.UploadServerCertificateInput{
		CertificateBody:       aws.String(d.Get("certificate_body").(string)),
		PrivateKey:            aws.String(d.Get("private_key").(string)),
		ServerCertificateName: aws.String(sslCertName),
	}

	if v, ok := d.GetOk("certificate_chain"); ok {
		createOpts.CertificateChain = aws.String(v.(string))
	}

	if v, ok := d.GetOk("path"); ok {
		createOpts.Path = aws.String(v.(string))
	}

	log.Printf("[DEBUG] Creating IAM Server Certificate with opts: %s", createOpts)
	resp, err := conn.UploadServerCertificate(createOpts)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			return fmt.Errorf("Error uploading server certificate, error: %s: %s", awsErr.Code(), awsErr.Message())
		}
		return fmt.Errorf("Error uploading server certificate, error: %s", err)
	}

	d.SetId(*resp.ServerCertificateMetadata.ServerCertificateId)
	d.Set("name", sslCertName)

	return resourceAwsIAMServerCertificateRead(d, meta)
}

func resourceAwsIAMServerCertificateRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).iamconn
	resp, err := conn.GetServerCertificate(&iam.GetServerCertificateInput{
		ServerCertificateName: aws.String(d.Get("name").(string)),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "NoSuchEntity" {
				log.Printf("[WARN] IAM Server Cert (%s) not found, removing from state", d.Id())
				d.SetId("")
				return nil
			}
			return fmt.Errorf("Error reading IAM Server Certificate: %s: %s", awsErr.Code(), awsErr.Message())
		}
		return fmt.Errorf("Error reading IAM Server Certificate: %s", err)
	}

	d.SetId(*resp.ServerCertificate.ServerCertificateMetadata.ServerCertificateId)

	// these values should always be present, and have a default if not set in
	// configuration, and so safe to reference with nil checks
	d.Set("certificate_body", normalizeCert(resp.ServerCertificate.CertificateBody))

	c := normalizeCert(resp.ServerCertificate.CertificateChain)
	if c != "" {
		d.Set("certificate_chain", c)
	}

	d.Set("path", resp.ServerCertificate.ServerCertificateMetadata.Path)
	d.Set("arn", resp.ServerCertificate.ServerCertificateMetadata.Arn)

	return nil
}

func resourceAwsIAMServerCertificateDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).iamconn
	log.Printf("[INFO] Deleting IAM Server Certificate: %s", d.Id())
	err := resource.Retry(15*time.Minute, func() *resource.RetryError {
		_, err := conn.DeleteServerCertificate(&iam.DeleteServerCertificateInput{
			ServerCertificateName: aws.String(d.Get("name").(string)),
		})

		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {
				if awsErr.Code() == "DeleteConflict" && strings.Contains(awsErr.Message(), "currently in use by arn") {
					currentlyInUseBy(awsErr.Message(), meta.(*AWSClient).elbconn)
					log.Printf("[WARN] Conflict deleting server certificate: %s, retrying", awsErr.Message())
					return resource.RetryableError(err)
				}
				if awsErr.Code() == "NoSuchEntity" {
					return nil
				}
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})

	return err
}

func resourceAwsIAMServerCertificateImport(
	d *schema.ResourceData, meta interface{}) ([]*schema.ResourceData, error) {
	d.Set("name", d.Id())
	// private_key can't be fetched from any API call
	return []*schema.ResourceData{d}, nil
}

func currentlyInUseBy(awsErr string, conn *elb.ELB) {
	r := regexp.MustCompile(`currently in use by ([a-z0-9:-]+)\/([a-z0-9-]+)\.`)
	matches := r.FindStringSubmatch(awsErr)
	if len(matches) > 0 {
		lbName := matches[2]
		describeElbOpts := &elb.DescribeLoadBalancersInput{
			LoadBalancerNames: []*string{aws.String(lbName)},
		}
		if _, err := conn.DescribeLoadBalancers(describeElbOpts); err != nil {
			if isAWSErr(err, "LoadBalancerNotFound", "") {
				log.Printf("[WARN] Load Balancer (%s) causing delete conflict not found", lbName)
			}
		}
	}
}

func normalizeCert(cert interface{}) string {
	if cert == nil || cert == (*string)(nil) {
		return ""
	}

	var rawCert string
	switch cert := cert.(type) {
	case string:
		rawCert = cert
	case *string:
		rawCert = *cert
	default:
		return ""
	}

	cleanVal := sha1.Sum(stripCR([]byte(strings.TrimSpace(rawCert))))
	return hex.EncodeToString(cleanVal[:])
}

// strip CRs from raw literals. Lifted from go/scanner/scanner.go
// See https://github.com/golang/go/blob/release-branch.go1.6/src/go/scanner/scanner.go#L479
func stripCR(b []byte) []byte {
	c := make([]byte, len(b))
	i := 0
	for _, ch := range b {
		if ch != '\r' {
			c[i] = ch
			i++
		}
	}
	return c[:i]
}
<commit_msg>Final retry for deleting IAM server cert<commit_after>package aws

import (
	"crypto/sha1"
	"encoding/hex"
	"fmt"
	"log"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
	"github.com/hashicorp/terraform/helper/validation"
)

func resourceAwsIAMServerCertificate() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsIAMServerCertificateCreate,
		Read:   resourceAwsIAMServerCertificateRead,
		Delete: resourceAwsIAMServerCertificateDelete,
		Importer: &schema.ResourceImporter{
			State: resourceAwsIAMServerCertificateImport,
		},

		Schema: map[string]*schema.Schema{
			"certificate_body": {
				Type:      schema.TypeString,
				Required:  true,
				ForceNew:  true,
				StateFunc: normalizeCert,
			},

			"certificate_chain": {
				Type:      schema.TypeString,
				Optional:  true,
				ForceNew:  true,
				StateFunc: normalizeCert,
			},

			"path": {
				Type:     schema.TypeString,
				Optional: true,
				Default:  "/",
				ForceNew: true,
			},

			"private_key": {
				Type:      schema.TypeString,
				Required:  true,
				ForceNew:  true,
				StateFunc: normalizeCert,
				Sensitive: true,
			},

			"name": {
				Type:          schema.TypeString,
				Optional:      true,
				Computed:      true,
				ForceNew:      true,
				ConflictsWith: []string{"name_prefix"},
				ValidateFunc:  validation.StringLenBetween(0, 128),
			},

			"name_prefix": {
				Type:          schema.TypeString,
				Optional:      true,
				ForceNew:      true,
				ConflictsWith: []string{"name"},
				ValidateFunc:  validation.StringLenBetween(0, 128-resource.UniqueIDSuffixLength),
			},

			"arn": {
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
		},
	}
}

func resourceAwsIAMServerCertificateCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).iamconn

	var sslCertName string
	if v, ok := d.GetOk("name"); ok {
		sslCertName = v.(string)
	} else if v, ok := d.GetOk("name_prefix"); ok {
		sslCertName = resource.PrefixedUniqueId(v.(string))
	} else {
		sslCertName = resource.UniqueId()
	}

	createOpts := &iam.UploadServerCertificateInput{
		CertificateBody:       aws.String(d.Get("certificate_body").(string)),
		PrivateKey:            aws.String(d.Get("private_key").(string)),
		ServerCertificateName: aws.String(sslCertName),
	}

	if v, ok := d.GetOk("certificate_chain"); ok {
		createOpts.CertificateChain = aws.String(v.(string))
	}

	if v, ok := d.GetOk("path"); ok {
		createOpts.Path = aws.String(v.(string))
	}

	log.Printf("[DEBUG] Creating IAM Server Certificate with opts: %s", createOpts)
	resp, err := conn.UploadServerCertificate(createOpts)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			return fmt.Errorf("Error uploading server certificate, error: %s: %s", awsErr.Code(), awsErr.Message())
		}
		return fmt.Errorf("Error uploading server certificate, error: %s", err)
	}

	d.SetId(*resp.ServerCertificateMetadata.ServerCertificateId)
	d.Set("name", sslCertName)

	return resourceAwsIAMServerCertificateRead(d, meta)
}

func resourceAwsIAMServerCertificateRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).iamconn
	resp, err := conn.GetServerCertificate(&iam.GetServerCertificateInput{
		ServerCertificateName: aws.String(d.Get("name").(string)),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "NoSuchEntity" {
				log.Printf("[WARN] IAM Server Cert (%s) not found, removing from state", d.Id())
				d.SetId("")
				return nil
			}
			return fmt.Errorf("Error reading IAM Server Certificate: %s: %s", awsErr.Code(), awsErr.Message())
		}
		return fmt.Errorf("Error reading IAM Server Certificate: %s", err)
	}

	d.SetId(*resp.ServerCertificate.ServerCertificateMetadata.ServerCertificateId)

	// these values should always be present, and have a default if not set in
	// configuration, and so safe to reference with nil checks
	d.Set("certificate_body", normalizeCert(resp.ServerCertificate.CertificateBody))

	c := normalizeCert(resp.ServerCertificate.CertificateChain)
	if c != "" {
		d.Set("certificate_chain", c)
	}

	d.Set("path", resp.ServerCertificate.ServerCertificateMetadata.Path)
	d.Set("arn", resp.ServerCertificate.ServerCertificateMetadata.Arn)

	return nil
}

func resourceAwsIAMServerCertificateDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).iamconn
	log.Printf("[INFO] Deleting IAM Server Certificate: %s", d.Id())
	err := resource.Retry(15*time.Minute, func() *resource.RetryError {
		_, err := conn.DeleteServerCertificate(&iam.DeleteServerCertificateInput{
			ServerCertificateName: aws.String(d.Get("name").(string)),
		})

		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {
				if awsErr.Code() == "DeleteConflict" && strings.Contains(awsErr.Message(), "currently in use by arn") {
					currentlyInUseBy(awsErr.Message(), meta.(*AWSClient).elbconn)
					log.Printf("[WARN] Conflict deleting server certificate: %s, retrying", awsErr.Message())
					return resource.RetryableError(err)
				}
				if awsErr.Code() == "NoSuchEntity" {
					return nil
				}
			}
			return resource.NonRetryableError(err)
		}
		return nil
	})

	if isResourceTimeoutError(err) {
		_, err = conn.DeleteServerCertificate(&iam.DeleteServerCertificateInput{
			ServerCertificateName: aws.String(d.Get("name").(string)),
		})
	}

	return err
}

func resourceAwsIAMServerCertificateImport(
	d *schema.ResourceData, meta interface{}) ([]*schema.ResourceData, error) {
	d.Set("name", d.Id())
	// private_key can't be fetched from any API call
	return []*schema.ResourceData{d}, nil
}

func currentlyInUseBy(awsErr string, conn *elb.ELB) {
	r := regexp.MustCompile(`currently in use by ([a-z0-9:-]+)\/([a-z0-9-]+)\.`)
	matches := r.FindStringSubmatch(awsErr)
	if len(matches) > 0 {
		lbName := matches[2]
		describeElbOpts := &elb.DescribeLoadBalancersInput{
			LoadBalancerNames: []*string{aws.String(lbName)},
		}
		if _, err := conn.DescribeLoadBalancers(describeElbOpts); err != nil {
			if isAWSErr(err, "LoadBalancerNotFound", "") {
				log.Printf("[WARN] Load Balancer (%s) causing delete conflict not found", lbName)
			}
		}
	}
}

func normalizeCert(cert interface{}) string {
	if cert == nil || cert == (*string)(nil) {
		return ""
	}

	var rawCert string
	switch cert := cert.(type) {
	case string:
		rawCert = cert
	case *string:
		rawCert = *cert
	default:
		return ""
	}

	cleanVal := sha1.Sum(stripCR([]byte(strings.TrimSpace(rawCert))))
	return hex.EncodeToString(cleanVal[:])
}

// strip CRs from raw literals. Lifted from go/scanner/scanner.go
// See https://github.com/golang/go/blob/release-branch.go1.6/src/go/scanner/scanner.go#L479
func stripCR(b []byte) []byte {
	c := make([]byte, len(b))
	i := 0
	for _, ch := range b {
		if ch != '\r' {
			c[i] = ch
			i++
		}
	}
	return c[:i]
}
<|endoftext|>