<commit_before>package utils

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"fmt"
	"log"
	"net/http"
	"os"
	"objects"
	"strconv"
	"html/template"
)

func HandlerRoot(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
    //ServeError(aResponseWriter, STR_MSG_NOTFOUND, STR_template_page_error_html)

    aRequest.ParseForm();

	cookie, errCookie := aRequest.Cookie(API_KEY_token)
	if errCookie != nil {
		log.Printf("handleRoot, Error reading cookie, error=%s", errCookie.Error())
		ServeLogin(aResponseWriter, STR_MSG_login);
		return
	}
	isTokenValid, userId := DbIsTokenValid(cookie.Value, nil)
	if errCookie == nil && !isTokenValid {
		ServeLogin(aResponseWriter, STR_MSG_login);
		return
	}

	user, errorUser := DbGetUserLoad(userId, nil);
	if errorUser != nil {
		log.Printf("errorUser=%s", errorUser.Error())
	}
	log.Printf("cookie.value=%s", cookie.Value)

	//Check if the file in the url path exists
	templateFile, err := template.ParseFiles(aRequest.URL.Path[1:]);
	if err != nil {
		ServeError(aResponseWriter, STR_MSG_404, STR_template_page_error_html);
	} else {
		if aRequest.URL.Path[1:] == "templates/Content.html" && user.Email != STR_EMPTY {
			err = templateFile.Execute(aResponseWriter, user);
		} else {
			err = templateFile.Execute(aResponseWriter, nil);
		}
		if err != nil {
			log.Printf("handleRoot, Error=", err.Error());
		}
	}
}

func HandlerEcho(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
	responseText := STR_EMPTY

	aRequest.ParseForm()

	body := aRequest.Form
	log.Printf("aRequest.Form=%s", body)

	bytesBody, err := ioutil.ReadAll(aRequest.Body)
	if(err != nil) {
		log.Printf("Error reading body, err=%s", err.Error())
	} else {
		log.Printf("bytesBody=%s", string(bytesBody))
		responseText = string(bytesBody)
	}

	headers := aRequest.Header
	for key, value := range headers {
		log.Printf("header=%s\n", key)
		fmt.Fprintf(aResponseWriter, "Header=%s\n", key)
		for idx, val := range value {
			log.Printf("idx=%d, value=%s", idx, val)
			fmt.Fprintf(aResponseWriter, "value=%s\n", val)
		}
	}

	fmt.Fprintf(aResponseWriter, "Method=%s\n", aRequest.Method)
	fmt.Fprintf(aResponseWriter, "%s\n", responseText)
}

func HandlerMessage(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
	aRequest.ParseForm()

	body := aRequest.Form
	log.Printf("aRequest.Form=%s", body)
	bytesBody, err := ioutil.ReadAll(aRequest.Body)
	if(err != nil) {
		log.Printf("Error reading body, err=%s", err.Error())
	} else {
		log.Printf("bytesBody=%s", string(bytesBody))
	}

	headerAuthentication := aRequest.Header.Get(STR_Authorization)
	log.Printf("headerAuthentication=%s", headerAuthentication)

	reportMessage := new(objects.ReportMessage)
	json.Unmarshal(bytesBody, reportMessage)
	log.Printf("report.Message=%s, report.Sequence=%d, report.Time=%d", reportMessage.Message, reportMessage.Sequence, reportMessage.Time)

	result := new(objects.Result)
	result.ErrorMessage = STR_EMPTY
	result.ResultCode = http.StatusOK
	ServeResult(aResponseWriter, result, STR_template_result)
}

func HandlerUploadImage(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
	requestMethod := aRequest.Method
	if requestMethod == STR_GET {
		result := new(objects.Result)
		result.ErrorMessage = STR_error
		result.ResultCode = http.StatusMethodNotAllowed
		ServeResult(aResponseWriter, result, STR_template_result)
	} else if requestMethod == STR_POST {
		//get message part
		errorParse := aRequest.ParseMultipartForm(8388608)
		if errorParse != nil {
			log.Printf("errorParse=%s", errorParse.Error())
		}
//		myform := aRequest.MultipartForm
//		valuesMap := myform.Value //map[string][]string
//		arrayMessage := valuesMap["message"]
//		log.Printf("arrayMessage=%d", len(arrayMessage))

		strMessage := aRequest.FormValue(API_KEY_message)
		log.Printf("strMessage=%s", strMessage)

		//get file part
		multipartFile, multipartFileHeader, err := aRequest.FormFile(API_KEY_image)
		if err != nil {
			log.Printf("Error getting file from FormFile, err=%s", err.Error())
			result := new(objects.Result)
			result.ErrorMessage = err.Error()
			result.ResultCode = http.StatusBadRequest
			ServeResult(aResponseWriter, result, STR_template_result)
			return
		}
		defer multipartFile.Close()

		imageFilePath := fmt.Sprintf(STR_img_filepathSave_template, multipartFileHeader.Filename)
		fileName := imageFilePath[0:(len(imageFilePath) - 4)]
		fileExstension := imageFilePath[(len(imageFilePath) - 4):len(imageFilePath)]
		fileNum := 0;
		var errorFileExists error
		_, errorFileExists = os.Stat(imageFilePath)
		for(!os.IsNotExist(errorFileExists)) {
			fileNum++
			imageFilePath = fileName + strconv.Itoa(fileNum) + fileExstension
			_, errorFileExists = os.Stat(imageFilePath)
		}
		log.Printf("imageFilePath=%s", imageFilePath)

		fileOut, errOut := os.Create(imageFilePath)
		if errOut != nil {
			log.Printf("Error creating fileOut, errOut=%s", errOut.Error())
			return
		}
		defer fileOut.Close()

		written, errWrite := io.Copy(fileOut, multipartFile)
		if errWrite != nil {
			log.Printf("Erro copying file, errWrite=%s", errWrite.Error())
			return
		}

		log.Printf("Bytes written=%d", written)

		result := new(objects.Result)
		result.ErrorMessage = STR_EMPTY
		result.ResultCode = http.StatusOK
		ServeResult(aResponseWriter, result, STR_template_result)
	}
}

func HandlerUploadFile(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
	requestMethod := aRequest.Method
	if requestMethod == STR_GET {
		result := new(objects.Result)
		result.ErrorMessage = STR_error
		result.ResultCode = http.StatusMethodNotAllowed
		ServeResult(aResponseWriter, result, STR_template_result)
	} else if requestMethod == STR_POST {
		multipartFile, multipartFileHeader, err := aRequest.FormFile(API_KEY_file)
		if err != nil {
			log.Printf("Error getting file from FormFile, err=%s", err.Error())
			result := new(objects.Result)
			result.ErrorMessage = err.Error()
			result.ResultCode = http.StatusBadRequest
			ServeResult(aResponseWriter, result, STR_template_result)
			return
		}
		defer multipartFile.Close()

		imageFilePath := fmt.Sprintf(STR_img_filepathSave_template, multipartFileHeader.Filename)
		fileName := imageFilePath[0:(len(imageFilePath) - 4)]
		fileExstension := imageFilePath[(len(imageFilePath) - 4):len(imageFilePath)]
		fileNum := 0;
		var errorFileExists error
		_, errorFileExists = os.Stat(imageFilePath)
		for(!os.IsNotExist(errorFileExists)) {
			fileNum++
			imageFilePath = fileName + strconv.Itoa(fileNum) + fileExstension
			_, errorFileExists = os.Stat(imageFilePath)
		}
		log.Printf("imageFilePath=%s", imageFilePath)

		fileOut, errOut := os.Create(imageFilePath)
		if errOut != nil {
			log.Printf("Error creating fileOut, errOut=%s", errOut.Error())
			return
		}
		defer fileOut.Close()

		written, errWrite := io.Copy(fileOut, multipartFile)
		if errWrite != nil {
			log.Printf("Erro copying file, errWrite=%s", errWrite.Error())
			return
		}

		log.Printf("Bytes written=%d", written)
		result := new(objects.Result)
		result.ErrorMessage = STR_EMPTY
		result.ResultCode = http.StatusOK
		ServeResult(aResponseWriter, result, STR_template_result)
	}
}

func HandlerLogin(responseWriter http.ResponseWriter, request *http.Request) {
	request.ParseForm();

	if request.Method == STR_GET {
		ServeLogin(responseWriter, STR_EMPTY);
	} else {
		var userName string = request.FormValue(API_KEY_username);
		var password string = request.FormValue(API_KEY_password);
		if userName == STR_EMPTY || password == STR_EMPTY {
			ServeLogin(responseWriter, "Please enter username and password");
			return;
		}

		var userId = -1;
		var errorUser error = nil
		userId, errorUser = DbGetUser(userName, password, nil)
		if errorUser != nil {
			log.Printf("handlerLogin, errorUser=%s", errorUser.Error())
		}
		if (userId > -1) {
			token := DbAddToken(userId, nil)
			AddCookie(responseWriter, token)
			http.Redirect(responseWriter, request, API_URL_Content, 301)
		} else {
			ServeLogin(responseWriter, "Wrong username or password");
		}
	}
}

func HandlerRegister(responseWriter http.ResponseWriter, request *http.Request) {
	request.ParseForm();

	if request.Method == STR_GET {
		ServeRegister(responseWriter, STR_EMPTY);
	} else {
		email := request.FormValue(API_KEY_email);
		password := request.FormValue(API_KEY_password);
		if (email == STR_EMPTY || password == STR_EMPTY) {
			ServeRegister(responseWriter, STR_MSG_register);
			return;
		}

		isUserExists, isUserAdded, errorUser := DbAddUser(email, password, nil);
		if errorUser != nil {
			log.Printf("handleRegister, errorUser=%s", errorUser.Error())
		}
		if isUserExists {
			ServeRegister(responseWriter, "Username is already taken.");
		} else if isUserAdded == false {
			ServeRegister(responseWriter, "Cannot create user.");
		} else {
			ServeLogin(responseWriter, STR_EMPTY);
		}
	}
}

/* Utils */
func getHeaderToken(aRequest *http.Request) string {
	headers := aRequest.Header
	tokens := headers["Token"]
	token := tokens[0]
	return token
}<commit_msg>1.Updating content handler.<commit_after>package utils

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"fmt"
	"log"
	"net/http"
	"os"
	"objects"
	"strconv"
	"html/template"
)

func HandlerRoot(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
    //ServeError(aResponseWriter, STR_MSG_NOTFOUND, STR_template_page_error_html)

    aRequest.ParseForm();

	cookie, errCookie := aRequest.Cookie(API_KEY_token)
	if errCookie != nil {
		log.Printf("handleRoot, Error reading cookie, error=%s", errCookie.Error())
		ServeLogin(aResponseWriter, STR_MSG_login);
		return
	}
	isTokenValid, userId := DbIsTokenValid(cookie.Value, nil)
	if errCookie == nil && !isTokenValid {
		ServeLogin(aResponseWriter, STR_MSG_login);
		return
	}

	user, errorUser := DbGetUserLoad(userId, nil);
	if errorUser != nil {
		log.Printf("errorUser=%s", errorUser.Error())
	}
	log.Printf("cookie.value=%s", cookie.Value)

	//Check if the file in the url path exists
	templateFile, err := template.ParseFiles(aRequest.URL.Path[1:]);
	if err != nil {
		ServeError(aResponseWriter, STR_MSG_404, STR_template_page_error_html);
	} else {
		AddCookie(aResponseWriter, cookie.Value)
		if aRequest.URL.Path[1:] == "templates/Content.html" && user.Email != STR_EMPTY {
			err = templateFile.Execute(aResponseWriter, user);
		} else {
			err = templateFile.Execute(aResponseWriter, nil);
		}
		if err != nil {
			log.Printf("handleRoot, Error=", err.Error());
		}
	}
}

func HandlerEcho(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
	responseText := STR_EMPTY

	aRequest.ParseForm()

	body := aRequest.Form
	log.Printf("aRequest.Form=%s", body)

	bytesBody, err := ioutil.ReadAll(aRequest.Body)
	if(err != nil) {
		log.Printf("Error reading body, err=%s", err.Error())
	} else {
		log.Printf("bytesBody=%s", string(bytesBody))
		responseText = string(bytesBody)
	}

	headers := aRequest.Header
	for key, value := range headers {
		log.Printf("header=%s\n", key)
		fmt.Fprintf(aResponseWriter, "Header=%s\n", key)
		for idx, val := range value {
			log.Printf("idx=%d, value=%s", idx, val)
			fmt.Fprintf(aResponseWriter, "value=%s\n", val)
		}
	}

	fmt.Fprintf(aResponseWriter, "Method=%s\n", aRequest.Method)
	fmt.Fprintf(aResponseWriter, "%s\n", responseText)
}

func HandlerMessage(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
	aRequest.ParseForm()

	body := aRequest.Form
	log.Printf("aRequest.Form=%s", body)
	bytesBody, err := ioutil.ReadAll(aRequest.Body)
	if(err != nil) {
		log.Printf("Error reading body, err=%s", err.Error())
	} else {
		log.Printf("bytesBody=%s", string(bytesBody))
	}

	headerAuthentication := aRequest.Header.Get(STR_Authorization)
	log.Printf("headerAuthentication=%s", headerAuthentication)

	reportMessage := new(objects.ReportMessage)
	json.Unmarshal(bytesBody, reportMessage)
	log.Printf("report.Message=%s, report.Sequence=%d, report.Time=%d", reportMessage.Message, reportMessage.Sequence, reportMessage.Time)

	result := new(objects.Result)
	result.ErrorMessage = STR_EMPTY
	result.ResultCode = http.StatusOK
	ServeResult(aResponseWriter, result, STR_template_result)
}

func HandlerUploadImage(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
	requestMethod := aRequest.Method
	if requestMethod == STR_GET {
		result := new(objects.Result)
		result.ErrorMessage = STR_error
		result.ResultCode = http.StatusMethodNotAllowed
		ServeResult(aResponseWriter, result, STR_template_result)
	} else if requestMethod == STR_POST {
		//get message part
		errorParse := aRequest.ParseMultipartForm(8388608)
		if errorParse != nil {
			log.Printf("errorParse=%s", errorParse.Error())
		}
//		myform := aRequest.MultipartForm
//		valuesMap := myform.Value //map[string][]string
//		arrayMessage := valuesMap["message"]
//		log.Printf("arrayMessage=%d", len(arrayMessage))

		strMessage := aRequest.FormValue(API_KEY_message)
		log.Printf("strMessage=%s", strMessage)

		//get file part
		multipartFile, multipartFileHeader, err := aRequest.FormFile(API_KEY_image)
		if err != nil {
			log.Printf("Error getting file from FormFile, err=%s", err.Error())
			result := new(objects.Result)
			result.ErrorMessage = err.Error()
			result.ResultCode = http.StatusBadRequest
			ServeResult(aResponseWriter, result, STR_template_result)
			return
		}
		defer multipartFile.Close()

		imageFilePath := fmt.Sprintf(STR_img_filepathSave_template, multipartFileHeader.Filename)
		fileName := imageFilePath[0:(len(imageFilePath) - 4)]
		fileExstension := imageFilePath[(len(imageFilePath) - 4):len(imageFilePath)]
		fileNum := 0;
		var errorFileExists error
		_, errorFileExists = os.Stat(imageFilePath)
		for(!os.IsNotExist(errorFileExists)) {
			fileNum++
			imageFilePath = fileName + strconv.Itoa(fileNum) + fileExstension
			_, errorFileExists = os.Stat(imageFilePath)
		}
		log.Printf("imageFilePath=%s", imageFilePath)

		fileOut, errOut := os.Create(imageFilePath)
		if errOut != nil {
			log.Printf("Error creating fileOut, errOut=%s", errOut.Error())
			return
		}
		defer fileOut.Close()

		written, errWrite := io.Copy(fileOut, multipartFile)
		if errWrite != nil {
			log.Printf("Erro copying file, errWrite=%s", errWrite.Error())
			return
		}

		log.Printf("Bytes written=%d", written)

		result := new(objects.Result)
		result.ErrorMessage = STR_EMPTY
		result.ResultCode = http.StatusOK
		ServeResult(aResponseWriter, result, STR_template_result)
	}
}

func HandlerUploadFile(aResponseWriter http.ResponseWriter, aRequest *http.Request) {
	requestMethod := aRequest.Method
	if requestMethod == STR_GET {
		result := new(objects.Result)
		result.ErrorMessage = STR_error
		result.ResultCode = http.StatusMethodNotAllowed
		ServeResult(aResponseWriter, result, STR_template_result)
	} else if requestMethod == STR_POST {
		multipartFile, multipartFileHeader, err := aRequest.FormFile(API_KEY_file)
		if err != nil {
			log.Printf("Error getting file from FormFile, err=%s", err.Error())
			result := new(objects.Result)
			result.ErrorMessage = err.Error()
			result.ResultCode = http.StatusBadRequest
			ServeResult(aResponseWriter, result, STR_template_result)
			return
		}
		defer multipartFile.Close()

		imageFilePath := fmt.Sprintf(STR_img_filepathSave_template, multipartFileHeader.Filename)
		fileName := imageFilePath[0:(len(imageFilePath) - 4)]
		fileExstension := imageFilePath[(len(imageFilePath) - 4):len(imageFilePath)]
		fileNum := 0;
		var errorFileExists error
		_, errorFileExists = os.Stat(imageFilePath)
		for(!os.IsNotExist(errorFileExists)) {
			fileNum++
			imageFilePath = fileName + strconv.Itoa(fileNum) + fileExstension
			_, errorFileExists = os.Stat(imageFilePath)
		}
		log.Printf("imageFilePath=%s", imageFilePath)

		fileOut, errOut := os.Create(imageFilePath)
		if errOut != nil {
			log.Printf("Error creating fileOut, errOut=%s", errOut.Error())
			return
		}
		defer fileOut.Close()

		written, errWrite := io.Copy(fileOut, multipartFile)
		if errWrite != nil {
			log.Printf("Erro copying file, errWrite=%s", errWrite.Error())
			return
		}

		log.Printf("Bytes written=%d", written)
		result := new(objects.Result)
		result.ErrorMessage = STR_EMPTY
		result.ResultCode = http.StatusOK
		ServeResult(aResponseWriter, result, STR_template_result)
	}
}

func HandlerLogin(responseWriter http.ResponseWriter, request *http.Request) {
	request.ParseForm();

	if request.Method == STR_GET {
		ServeLogin(responseWriter, STR_EMPTY);
	} else {
		var userName string = request.FormValue(API_KEY_username);
		var password string = request.FormValue(API_KEY_password);
		if userName == STR_EMPTY || password == STR_EMPTY {
			ServeLogin(responseWriter, "Please enter username and password");
			return;
		}

		var userId = -1;
		var errorUser error = nil
		userId, errorUser = DbGetUser(userName, password, nil)
		if errorUser != nil {
			log.Printf("handlerLogin, errorUser=%s", errorUser.Error())
		}
		if (userId > -1) {
			token := DbAddToken(userId, nil)
			AddCookie(responseWriter, token)
			http.Redirect(responseWriter, request, API_URL_Content, 301)
		} else {
			ServeLogin(responseWriter, "Wrong username or password");
		}
	}
}

func HandlerRegister(responseWriter http.ResponseWriter, request *http.Request) {
	request.ParseForm();

	if request.Method == STR_GET {
		ServeRegister(responseWriter, STR_EMPTY);
	} else {
		email := request.FormValue(API_KEY_email);
		password := request.FormValue(API_KEY_password);
		if (email == STR_EMPTY || password == STR_EMPTY) {
			ServeRegister(responseWriter, STR_MSG_register);
			return;
		}

		isUserExists, isUserAdded, errorUser := DbAddUser(email, password, nil);
		if errorUser != nil {
			log.Printf("handleRegister, errorUser=%s", errorUser.Error())
		}
		if isUserExists {
			ServeRegister(responseWriter, "Username is already taken.");
		} else if isUserAdded == false {
			ServeRegister(responseWriter, "Cannot create user.");
		} else {
			ServeLogin(responseWriter, STR_EMPTY);
		}
	}
}

/* Utils */
func getHeaderToken(aRequest *http.Request) string {
	headers := aRequest.Header
	tokens := headers["Token"]
	token := tokens[0]
	return token
}<|endoftext|>


<commit_before>package restful

// Copyright 2013 Ernest Micklei. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

import (
    "encoding/json"
    "encoding/xml"
    "net/http"
    "reflect"
    "strings"
)

// Added by apoorva moghey to generate custom response format
type Rejoinder struct {
    Success     bool   `json:"success"`
    ErrorCode   int    `json:"error_code,omitempty"`
    Message     string `json:"message"`
    TotalPage   int    `json:"total_page,omitempty"`
    CurrentPage int    `json:"current_page,omitempty"`
    Total       int    `json:"total,omitempty"`
}

// Added by apoorva moghey to generate custom response format
type Fallacy struct {
    Success   bool   `json:"success"`
    ErrorCode int    `json:"error_code,omitempty"`
    Message   string `json:"message"`
}

// DEPRECATED, use DefaultResponseContentType(mime)
var DefaultResponseMimeType string

//PrettyPrintResponses controls the indentation feature of XML and JSON
//serialization in the response methods WriteEntity, WriteAsJson, and
//WriteAsXml.
var PrettyPrintResponses = true

// Response is a wrapper on the actual http ResponseWriter
// It provides several convenience methods to prepare and write response content.
type Response struct {
    http.ResponseWriter
    requestAccept string   // mime-type what the Http Request says it wants to receive
    routeProduces []string // mime-types what the Route says it can produce
    statusCode    int      // HTTP status code that has been written explicity (if zero then net/http has written 200)
    contentLength int      // number of bytes written for the response body
    prettyPrint   bool     // controls the indentation feature of XML and JSON serialization. It is initialized using var PrettyPrintResponses.
    Data          interface{}
    Source        string
}

// Creates a new response based on a http ResponseWriter.
func NewResponse(httpWriter http.ResponseWriter) *Response {
    return &Response{
        httpWriter,
        "",
        []string{},
        http.StatusOK,
        0,
        PrettyPrintResponses, nil, ""} // empty content-types
}

// If Accept header matching fails, fall back to this type, otherwise
// a "406: Not Acceptable" response is returned.
// Valid values are restful.MIME_JSON and restful.MIME_XML
// Example:
// 	restful.DefaultResponseContentType(restful.MIME_JSON)
func DefaultResponseContentType(mime string) {
    DefaultResponseMimeType = mime
}

// InternalServerError writes the StatusInternalServerError header.
// DEPRECATED, use WriteErrorString(http.StatusInternalServerError,reason)
func (r Response) InternalServerError() Response {
    r.WriteHeader(http.StatusInternalServerError)
    return r
}

// PrettyPrint changes whether this response must produce pretty (line-by-line, indented) JSON or XML output.
func (r *Response) PrettyPrint(bePretty bool) {
    r.prettyPrint = bePretty
}

// AddHeader is a shortcut for .Header().Add(header,value)
func (r Response) AddHeader(header string, value string) Response {
    r.Header().Add(header, value)
    return r
}

// SetRequestAccepts tells the response what Mime-type(s) the HTTP request said it wants to accept. Exposed for testing.
func (r *Response) SetRequestAccepts(mime string) {
    r.requestAccept = mime
}

// WriteEntity marshals the value using the representation denoted by the Accept Header (XML or JSON)
// If no Accept header is specified (or */*) then return the Content-Type as specified by the first in the Route.Produces.
// If an Accept header is specified then return the Content-Type as specified by the first in the Route.Produces that is matched with the Accept header.
// If the value is nil then nothing is written. You may want to call WriteHeader(http.StatusNotFound) instead.
// Current implementation ignores any q-parameters in the Accept Header.
func (r *Response) WriteEntity(value interface{}) error {
    if reflect.TypeOf(reflect.TypeOf(value).Field(0)).NumField() != 7 {
        panic("Output must extend Rejoinder struct.")
    }

    if r.Source == "external" {
        r.Data = value
        r.Source = ""
        return nil
    }

    if value == nil { // do not write a nil representation
        return nil
    }
    for _, qualifiedMime := range strings.Split(r.requestAccept, ",") {
        mime := strings.Trim(strings.Split(qualifiedMime, ";")[0], " ")
        if 0 == len(mime) || mime == "*/*" {
            for _, each := range r.routeProduces {
                if MIME_JSON == each {
                    return r.WriteAsJson(value)
                }
                if MIME_XML == each {
                    return r.WriteAsXml(value)
                }
            }
        } else { // mime is not blank; see if we have a match in Produces
            for _, each := range r.routeProduces {
                if mime == each {
                    if MIME_JSON == each {
                        return r.WriteAsJson(value)
                    }
                    if MIME_XML == each {
                        return r.WriteAsXml(value)
                    }
                }
            }
        }
    }
    if DefaultResponseMimeType == MIME_JSON {
        return r.WriteAsJson(value)
    } else if DefaultResponseMimeType == MIME_XML {
        return r.WriteAsXml(value)
    } else {
        if trace {
            traceLogger.Printf("mismatch in mime-types and no defaults; (http)Accept=%v,(route)Produces=%v\n", r.requestAccept, r.routeProduces)
        }
        r.WriteHeader(http.StatusNotAcceptable) // for recording only
        r.ResponseWriter.WriteHeader(http.StatusNotAcceptable)
        if _, err := r.Write([]byte("406: Not Acceptable")); err != nil {
            return err
        }
    }
    return nil
}

// WriteAsXml is a convenience method for writing a value in xml (requires Xml tags on the value)
func (r *Response) WriteAsXml(value interface{}) error {
    var output []byte
    var err error

    if value == nil { // do not write a nil representation
        return nil
    }
    if r.prettyPrint {
        output, err = xml.MarshalIndent(value, " ", " ")
    } else {
        output, err = xml.Marshal(value)
    }

    if err != nil {
        return r.WriteError(http.StatusInternalServerError, err)
    }
    r.Header().Set(HEADER_ContentType, MIME_XML)
    if r.statusCode > 0 { // a WriteHeader was intercepted
        r.ResponseWriter.WriteHeader(r.statusCode)
    }
    _, err = r.Write([]byte(xml.Header))
    if err != nil {
        return err
    }
    if _, err = r.Write(output); err != nil {
        return err
    }
    return nil
}

// WriteAsJson is a convenience method for writing a value in json
func (r *Response) WriteAsJson(value interface{}) error {
    var output []byte
    var err error

    if value == nil { // do not write a nil representation
        return nil
    }
    if r.prettyPrint {
        output, err = json.MarshalIndent(value, " ", " ")
    } else {
        output, err = json.Marshal(value)
    }

    if err != nil {
        return r.WriteErrorString(http.StatusInternalServerError, err.Error())
    }
    r.Header().Set(HEADER_ContentType, MIME_JSON)
    if r.statusCode > 0 { // a WriteHeader was intercepted
        r.ResponseWriter.WriteHeader(r.statusCode)
    }
    if _, err = r.Write(output); err != nil {
        return err
    }
    return nil
}

// WriteError write the http status and the error string on the response.
func (r *Response) WriteError(httpStatus int, err error) error {
    return r.WriteErrorString(httpStatus, err.Error())
}

// WriteServiceError is a convenience method for a responding with a ServiceError and a status
func (r *Response) WriteServiceError(httpStatus int, err ServiceError) error {
    r.WriteHeader(httpStatus) // for recording only
    return r.WriteEntity(err)
}

// WriteErrorString is a convenience method for an error status with the actual error
func (r *Response) WriteErrorString(status int, errorReason string) error {
    r.statusCode = status // for recording only
    r.ResponseWriter.WriteHeader(status)
    if _, err := r.Write([]byte(errorReason)); err != nil {
        return err
    }
    return nil
}

// WriteHeader is overridden to remember the Status Code that has been written.
// Note that using this method, the status value is only written when
// - calling WriteEntity,
// - or directly calling WriteAsXml or WriteAsJson,
// - or if the status is one for which no response is allowed (i.e.,
//   204 (http.StatusNoContent) or 304 (http.StatusNotModified))
func (r *Response) WriteHeader(httpStatus int) {
    r.statusCode = httpStatus
    // if 204 then WriteEntity will not be called so we need to pass this code
    if http.StatusNoContent == httpStatus ||
        http.StatusNotModified == httpStatus {
        r.ResponseWriter.WriteHeader(httpStatus)
    }
}

// StatusCode returns the code that has been written using WriteHeader.
func (r Response) StatusCode() int {
    if 0 == r.statusCode {
        // no status code has been written yet; assume OK
        return http.StatusOK
    }
    return r.statusCode
}

// Write writes the data to the connection as part of an HTTP reply.
// Write is part of http.ResponseWriter interface.
func (r *Response) Write(bytes []byte) (int, error) {
    written, err := r.ResponseWriter.Write(bytes)
    r.contentLength += written
    return written, err
}

// ContentLength returns the number of bytes written for the response content.
// Note that this value is only correct if all data is written through the Response using its Write* methods.
// Data written directly using the underlying http.ResponseWriter is not accounted for.
func (r Response) ContentLength() int {
    return r.contentLength
}

// CloseNotify is part of http.CloseNotifier interface
func (r Response) CloseNotify() <-chan bool {
    return r.ResponseWriter.(http.CloseNotifier).CloseNotify()
}

func (r *Response) RemoveSource() {
    r.Source = ""
}

func (r *Response) WriteResp(value interface{}) error {
    if reflect.TypeOf(reflect.TypeOf(value).Field(0)).NumField() != 7 {
        panic("Output must extend Rejoinder struct.")
    }

    if value == nil { // do not write a nil representation
        return nil
    }
    for _, qualifiedMime := range strings.Split(r.requestAccept, ",") {
        mime := strings.Trim(strings.Split(qualifiedMime, ";")[0], " ")
        if 0 == len(mime) || mime == "*/*" {
            for _, each := range r.routeProduces {
                if MIME_JSON == each {
                    return r.WriteAsJson(value)
                }
                if MIME_XML == each {
                    return r.WriteAsXml(value)
                }
            }
        } else { // mime is not blank; see if we have a match in Produces
            for _, each := range r.routeProduces {
                if mime == each {
                    if MIME_JSON == each {
                        return r.WriteAsJson(value)
                    }
                    if MIME_XML == each {
                        return r.WriteAsXml(value)
                    }
                }
            }
        }
    }
    if DefaultResponseMimeType == MIME_JSON {
        return r.WriteAsJson(value)
    } else if DefaultResponseMimeType == MIME_XML {
        return r.WriteAsXml(value)
    } else {
        if trace {
            traceLogger.Printf("mismatch in mime-types and no defaults; (http)Accept=%v,(route)Produces=%v\n", r.requestAccept, r.routeProduces)
        }
        r.WriteHeader(http.StatusNotAcceptable) // for recording only
        r.ResponseWriter.WriteHeader(http.StatusNotAcceptable)
        if _, err := r.Write([]byte("406: Not Acceptable")); err != nil {
            return err
        }
    }
    return nil
}
<commit_msg>add common response method<commit_after>package restful

// Copyright 2013 Ernest Micklei. All rights reserved.
// Use of this source code is governed by a license
// that can be found in the LICENSE file.

import (
    "encoding/json"
    "encoding/xml"
    "net/http"
    "reflect"
    "strings"
)

// Added by apoorva moghey to generate custom response format
type Rejoinder struct {
    Success     bool   `json:"success"`
    ErrorCode   int    `json:"error_code,omitempty"`
    Message     string `json:"message"`
    Total       int    `json:"total,omitempty"`
    TotalPage   int    `json:"total_page,omitempty"`
    CurrentPage int    `json:"current_page,omitempty"`
}

// Added by apoorva moghey to generate custom response format
type Fallacy struct {
    Success   bool   `json:"success"`
    ErrorCode int    `json:"error_code,omitempty"`
    Message   string `json:"message"`
}

type Rep struct {
    Rejoinder
    Data interface{} `json:"data,omitempty"`
}

type RepError struct {
    Fallacy
    Data interface{} `json:"data,omitempty"`
}

// DEPRECATED, use DefaultResponseContentType(mime)
var DefaultResponseMimeType string

//PrettyPrintResponses controls the indentation feature of XML and JSON
//serialization in the response methods WriteEntity, WriteAsJson, and
//WriteAsXml.
var PrettyPrintResponses = true

// Response is a wrapper on the actual http ResponseWriter
// It provides several convenience methods to prepare and write response content.
type Response struct {
    http.ResponseWriter
    requestAccept string   // mime-type what the Http Request says it wants to receive
    routeProduces []string // mime-types what the Route says it can produce
    statusCode    int      // HTTP status code that has been written explicity (if zero then net/http has written 200)
    contentLength int      // number of bytes written for the response body
    prettyPrint   bool     // controls the indentation feature of XML and JSON serialization. It is initialized using var PrettyPrintResponses.
    Data          interface{}
    Source        string
}

// Creates a new response based on a http ResponseWriter.
func NewResponse(httpWriter http.ResponseWriter) *Response {
    return &Response{
        httpWriter,
        "",
        []string{},
        http.StatusOK,
        0,
        PrettyPrintResponses, nil, ""} // empty content-types
}

// If Accept header matching fails, fall back to this type, otherwise
// a "406: Not Acceptable" response is returned.
// Valid values are restful.MIME_JSON and restful.MIME_XML
// Example:
// 	restful.DefaultResponseContentType(restful.MIME_JSON)
func DefaultResponseContentType(mime string) {
    DefaultResponseMimeType = mime
}

// InternalServerError writes the StatusInternalServerError header.
// DEPRECATED, use WriteErrorString(http.StatusInternalServerError,reason)
func (r Response) InternalServerError() Response {
    r.WriteHeader(http.StatusInternalServerError)
    return r
}

// PrettyPrint changes whether this response must produce pretty (line-by-line, indented) JSON or XML output.
func (r *Response) PrettyPrint(bePretty bool) {
    r.prettyPrint = bePretty
}

// AddHeader is a shortcut for .Header().Add(header,value)
func (r Response) AddHeader(header string, value string) Response {
    r.Header().Add(header, value)
    return r
}

// SetRequestAccepts tells the response what Mime-type(s) the HTTP request said it wants to accept. Exposed for testing.
func (r *Response) SetRequestAccepts(mime string) {
    r.requestAccept = mime
}

// WriteEntity marshals the value using the representation denoted by the Accept Header (XML or JSON)
// If no Accept header is specified (or */*) then return the Content-Type as specified by the first in the Route.Produces.
// If an Accept header is specified then return the Content-Type as specified by the first in the Route.Produces that is matched with the Accept header.
// If the value is nil then nothing is written. You may want to call WriteHeader(http.StatusNotFound) instead.
// Current implementation ignores any q-parameters in the Accept Header.
func (r *Response) WriteEntity(value interface{}) error {
    if value == nil { // do not write a nil representation
        return nil
    }
    for _, qualifiedMime := range strings.Split(r.requestAccept, ",") {
        mime := strings.Trim(strings.Split(qualifiedMime, ";")[0], " ")
        if 0 == len(mime) || mime == "*/*" {
            for _, each := range r.routeProduces {
                if MIME_JSON == each {
                    return r.WriteAsJson(value)
                }
                if MIME_XML == each {
                    return r.WriteAsXml(value)
                }
            }
        } else { // mime is not blank; see if we have a match in Produces
            for _, each := range r.routeProduces {
                if mime == each {
                    if MIME_JSON == each {
                        return r.WriteAsJson(value)
                    }
                    if MIME_XML == each {
                        return r.WriteAsXml(value)
                    }
                }
            }
        }
    }
    if DefaultResponseMimeType == MIME_JSON {
        return r.WriteAsJson(value)
    } else if DefaultResponseMimeType == MIME_XML {
        return r.WriteAsXml(value)
    } else {
        if trace {
            traceLogger.Printf("mismatch in mime-types and no defaults; (http)Accept=%v,(route)Produces=%v\n", r.requestAccept, r.routeProduces)
        }
        r.WriteHeader(http.StatusNotAcceptable) // for recording only
        r.ResponseWriter.WriteHeader(http.StatusNotAcceptable)
        if _, err := r.Write([]byte("406: Not Acceptable")); err != nil {
            return err
        }
    }
    return nil
}

// WriteAsXml is a convenience method for writing a value in xml (requires Xml tags on the value)
func (r *Response) WriteAsXml(value interface{}) error {
    var output []byte
    var err error

    if value == nil { // do not write a nil representation
        return nil
    }
    if r.prettyPrint {
        output, err = xml.MarshalIndent(value, " ", " ")
    } else {
        output, err = xml.Marshal(value)
    }

    if err != nil {
        return r.WriteError(http.StatusInternalServerError, err)
    }
    r.Header().Set(HEADER_ContentType, MIME_XML)
    if r.statusCode > 0 { // a WriteHeader was intercepted
        r.ResponseWriter.WriteHeader(r.statusCode)
    }
    _, err = r.Write([]byte(xml.Header))
    if err != nil {
        return err
    }
    if _, err = r.Write(output); err != nil {
        return err
    }
    return nil
}

// WriteAsJson is a convenience method for writing a value in json
func (r *Response) WriteAsJson(value interface{}) error {
    var output []byte
    var err error

    if value == nil { // do not write a nil representation
        return nil
    }
    if r.prettyPrint {
        output, err = json.MarshalIndent(value, " ", " ")
    } else {
        output, err = json.Marshal(value)
    }

    if err != nil {
        return r.WriteErrorString(http.StatusInternalServerError, err.Error())
    }
    r.Header().Set(HEADER_ContentType, MIME_JSON)
    if r.statusCode > 0 { // a WriteHeader was intercepted
        r.ResponseWriter.WriteHeader(r.statusCode)
    }
    if _, err = r.Write(output); err != nil {
        return err
    }
    return nil
}

// WriteError write the http status and the error string on the response.
func (r *Response) WriteError(httpStatus int, err error) error {
    return r.WriteErrorString(httpStatus, err.Error())
}

// WriteServiceError is a convenience method for a responding with a ServiceError and a status
func (r *Response) WriteServiceError(httpStatus int, err ServiceError) error {
    r.WriteHeader(httpStatus) // for recording only
    return r.WriteEntity(err)
}

// WriteErrorString is a convenience method for an error status with the actual error
func (r *Response) WriteErrorString(status int, errorReason string) error {
    r.statusCode = status // for recording only
    r.ResponseWriter.WriteHeader(status)
    if _, err := r.Write([]byte(errorReason)); err != nil {
        return err
    }
    return nil
}

// WriteHeader is overridden to remember the Status Code that has been written.
// Note that using this method, the status value is only written when
// - calling WriteEntity,
// - or directly calling WriteAsXml or WriteAsJson,
// - or if the status is one for which no response is allowed (i.e.,
//   204 (http.StatusNoContent) or 304 (http.StatusNotModified))
func (r *Response) WriteHeader(httpStatus int) {
    r.statusCode = httpStatus
    // if 204 then WriteEntity will not be called so we need to pass this code
    if http.StatusNoContent == httpStatus ||
        http.StatusNotModified == httpStatus {
        r.ResponseWriter.WriteHeader(httpStatus)
    }
}

// StatusCode returns the code that has been written using WriteHeader.
func (r Response) StatusCode() int {
    if 0 == r.statusCode {
        // no status code has been written yet; assume OK
        return http.StatusOK
    }
    return r.statusCode
}

// Write writes the data to the connection as part of an HTTP reply.
// Write is part of http.ResponseWriter interface.
func (r *Response) Write(bytes []byte) (int, error) {
    written, err := r.ResponseWriter.Write(bytes)
    r.contentLength += written
    return written, err
}

// ContentLength returns the number of bytes written for the response content.
// Note that this value is only correct if all data is written through the Response using its Write* methods.
// Data written directly using the underlying http.ResponseWriter is not accounted for.
func (r Response) ContentLength() int {
    return r.contentLength
}

// CloseNotify is part of http.CloseNotifier interface
func (r Response) CloseNotify() <-chan bool {
    return r.ResponseWriter.(http.CloseNotifier).CloseNotify()
}

func (r *Response) RemoveSource() {
    r.Source = ""
}

func (r *Response) Reply(data interface{}, message string,  meta ...int) {
    resp := Rep {
        Rejoinder {
	        Success: true,
	        Message: message,
	        Total : len(data),
	        TotalPage: meta[0],
	        CurrentPage:meta[1],
        },data}
    r.WriteEntity(&resp)
}

func (r *Response) ReplyError(data interface{}, message string, errorCode int) {
	resp := RepError {
	    Fallacy {
	        Success : false,
	        ErrorCode : errorCode,
	        Message : message,
	    },data}
	r.WriteEntity(&resp)
}

func len(v interface{}) int {
	if v == nil {
		return 0
	}
	switch reflect.ValueOf(v).Kind() {
	case reflect.Ptr:
	case reflect.Struct:
		return 1
	default:
		return reflect.ValueOf(v).Len()
	}
	return 1
}<|endoftext|>


<commit_before>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package table

/*  Filename:    testing.go
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Sat Dec 10 15:09:48 PST 2011
 *  Description:
 */

import (
//"testing"
)

type testingT struct {
	name string
	t    Testing
}

func newTestingT(name string, t Testing) *testingT { return &testingT{name, t} }
func (t *testingT) dup() (cp *testingT)            { cp = new(testingT); *cp = *t; return }
func (t *testingT) sub(name string) (s *testingT)  { s = newTestingT(name, t); return }

func (t *testingT) msg(v ...interface{}) string            {
	m := sprint(v...)
	if t.name != "" {
		m = msg(t.name, m)
	}
	return m
}
func (t *testingT) errmsg(typ string, v ...interface{}) string {
	name, m := t.name, sprint(v...)
	if Verbose {
		prefix := typ
		if name != "" {
			prefix = sprintf("%s %s", t.name, typ)
		}
		name = msg(name, prefix)
	}
	if name != "" {
		m = msg(name, m)
	}
	return m
}
func (t *testingT) msgf(f string, v ...interface{}) string { return msg(t.name, sprintf(f, v...)) }

func (t *testingT) Fail()                                     { t.t.Fail() }
func (t *testingT) FailNow()                                  { t.t.FailNow() }
func (t *testingT) Failed() bool                              { return t.t.Failed() }
func (t *testingT) log(args ...interface{})                   { t.t.Log(sprint(args...)) }
func (t *testingT) error(args ...interface{})                 { t.t.Error(sprint(args...)) }
func (t *testingT) fatal(args ...interface{})                 { t.t.Fatal(sprint(args...)) }
func (t *testingT) Log(args ...interface{})                   { t.log(t.msg(args...)) }
func (t *testingT) Error(args ...interface{})                 { t.error(t.errmsg("error", args...)) }
func (t *testingT) Fatal(args ...interface{})                 { t.fatal(t.errmsg("fatal", args...)) }
func (t *testingT) Logf(format string, args ...interface{})   { t.log(t.msgf(format, args...)) }
func (t *testingT) Errorf(format string, args ...interface{}) { t.error(t.msgf(format, args...)) }
func (t *testingT) Fatalf(format string, args ...interface{}) { t.fatal(t.msgf(format, args...)) }

type Testing interface {
	Error(args ...interface{})
	Errorf(format string, args ...interface{})
	Fail()
	FailNow()
	Failed() bool
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Log(args ...interface{})
	Logf(format string, args ...interface{})
}
<commit_msg>Fix another bug with verbose error messages.<commit_after>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package table

/*  Filename:    testing.go
 *  Author:      Bryan Matsuo <bmatsuo@soe.ucsc.edu>
 *  Created:     Sat Dec 10 15:09:48 PST 2011
 *  Description:
 */

import (
//"testing"
)

type testingT struct {
	name string
	t    Testing
}

func newTestingT(name string, t Testing) *testingT { return &testingT{name, t} }
func (t *testingT) dup() (cp *testingT)            { cp = new(testingT); *cp = *t; return }
func (t *testingT) sub(name string) (s *testingT)  { s = newTestingT(name, t); return }

func (t *testingT) msg(v ...interface{}) string            {
	m := sprint(v...)
	if t.name != "" {
		m = msg(t.name, m)
	}
	return m
}
func (t *testingT) errmsg(typ string, v ...interface{}) string {
	name, m := t.name, sprint(v...)
	if Verbose {
		prefix := typ
		if name != "" {
			prefix = sprintf("%s %s", t.name, typ)
		}
		name = msg(name, prefix)
	}
	if name != "" {
		m = msg(name, m)
	}
	return m
}
func (t *testingT) msgf(f string, v ...interface{}) string { return msg(t.name, sprintf(f, v...)) }

func (t *testingT) Fail()                                     { t.t.Fail() }
func (t *testingT) FailNow()                                  { t.t.FailNow() }
func (t *testingT) Failed() bool                              { return t.t.Failed() }
func (t *testingT) log(args ...interface{})                   { t.t.Log(sprint(args...)) }
func (t *testingT) error(args ...interface{})                 { t.t.Error(sprint(args...)) }
func (t *testingT) fatal(args ...interface{})                 { t.t.Fatal(sprint(args...)) }
func (t *testingT) Log(args ...interface{})                   { t.log(t.msg(args...)) }
func (t *testingT) Error(args ...interface{})                 { t.error(t.errmsg("error", args...)) }
func (t *testingT) Fatal(args ...interface{})                 { t.fatal(t.errmsg("fatal", args...)) }
func (t *testingT) Logf(format string, args ...interface{})   { t.Log(sprintf(format, args...)) }
func (t *testingT) Errorf(format string, args ...interface{}) { t.Error(sprintf(format, args...)) }
func (t *testingT) Fatalf(format string, args ...interface{}) { t.Fatal(sprintf(format, args...)) }

type Testing interface {
	Error(args ...interface{})
	Errorf(format string, args ...interface{})
	Fail()
	FailNow()
	Failed() bool
	Fatal(args ...interface{})
	Fatalf(format string, args ...interface{})
	Log(args ...interface{})
	Logf(format string, args ...interface{})
}
<|endoftext|>


<commit_before>package chunker

import (
	"errors"
	"io"
	"sync"
)

const (
	kiB = 1024
	miB = 1024 * kiB

	// WindowSize is the size of the sliding window.
	windowSize = 64

	// MinSize is the default minimal size of a chunk.
	MinSize = 512 * kiB
	// MaxSize is the default maximal size of a chunk.
	MaxSize = 8 * miB

	chunkerBufSize = 512 * kiB
)

type tables struct {
	out [256]Pol
	mod [256]Pol
}

// cache precomputed tables, these are read-only anyway
var cache struct {
	entries map[Pol]tables
	sync.Mutex
}

func init() {
	cache.entries = make(map[Pol]tables)
}

// Chunk is one content-dependent chunk of bytes whose end was cut when the
// Rabin Fingerprint had the value stored in Cut.
type Chunk struct {
	Start  uint
	Length uint
	Cut    uint64
	Data   []byte
}

type chunkerState struct {
	window [windowSize]byte
	wpos   int

	buf  []byte
	bpos uint
	bmax uint

	start uint
	count uint
	pos   uint

	pre uint // wait for this many bytes before start calculating an new chunk

	digest uint64
}

type chunkerConfig struct {
	MinSize, MaxSize uint

	pol               Pol
	polShift          uint
	tables            tables
	tablesInitialized bool
	splitmask         uint64

	rd     io.Reader
	closed bool
}

// Chunker splits content with Rabin Fingerprints.
type Chunker struct {
	chunkerConfig
	chunkerState
}

// SetAverageBits allows to control the frequency of chunk discovery:
// the lower averageBits, the higher amount of chunks will be identified.
// The default value is 20 bits, so chunks will be of 1MiB size on average.
func (c *Chunker) SetAverageBits(averageBits int) {
	c.splitmask = (1 << uint64(averageBits)) - 1
}

// New returns a new Chunker based on polynomial p that reads from rd.
func New(rd io.Reader, pol Pol) *Chunker {
	return NewWithBoundaries(rd, pol, MinSize, MaxSize)
}

// NewWithBoundaries returns a new Chunker based on polynomial p that reads from
// rd and custom min and max size boundaries.
func NewWithBoundaries(rd io.Reader, pol Pol, min, max uint) *Chunker {
	c := &Chunker{
		chunkerState: chunkerState{
			buf: make([]byte, chunkerBufSize),
		},
		chunkerConfig: chunkerConfig{
			pol:       pol,
			rd:        rd,
			MinSize:   min,
			MaxSize:   max,
			splitmask: (1 << 20) - 1, // aim to create chunks of 20 bits or about 1MiB on average.
		},
	}

	c.reset()

	return c
}

// Reset reinitializes the chunker with a new reader and polynomial.
func (c *Chunker) Reset(rd io.Reader, pol Pol) {
	c.ResetWithBoundaries(rd, pol, MinSize, MaxSize)
}

// ResetWithBoundaries reinitializes the chunker with a new reader, polynomial
// and custom min and max size boundaries.
func (c *Chunker) ResetWithBoundaries(rd io.Reader, pol Pol, min, max uint) {
	*c = Chunker{
		chunkerState: chunkerState{
			buf: c.buf,
		},
		chunkerConfig: chunkerConfig{
			pol:       pol,
			rd:        rd,
			MinSize:   min,
			MaxSize:   max,
			splitmask: (1 << 20) - 1,
		},
	}

	c.reset()
}

func (c *Chunker) reset() {
	c.polShift = uint(c.pol.Deg() - 8)
	c.fillTables()

	for i := 0; i < windowSize; i++ {
		c.window[i] = 0
	}

	c.closed = false
	c.digest = 0
	c.wpos = 0
	c.count = 0
	c.digest = c.slide(c.digest, 1)
	c.start = c.pos

	// do not start a new chunk unless at least MinSize bytes have been read
	c.pre = c.MinSize - windowSize
}

// fillTables calculates out_table and mod_table for optimization. This
// implementation uses a cache in the global variable cache.
func (c *Chunker) fillTables() {
	// if polynomial hasn't been specified, do not compute anything for now
	if c.pol == 0 {
		return
	}

	c.tablesInitialized = true

	// test if the tables are cached for this polynomial
	cache.Lock()
	defer cache.Unlock()
	if t, ok := cache.entries[c.pol]; ok {
		c.tables = t
		return
	}

	// calculate table for sliding out bytes. The byte to slide out is used as
	// the index for the table, the value contains the following:
	// out_table[b] = Hash(b || 0 ||        ...        || 0)
	//                          \ windowsize-1 zero bytes /
	// To slide out byte b_0 for window size w with known hash
	// H := H(b_0 || ... || b_w), it is sufficient to add out_table[b_0]:
	//    H(b_0 || ... || b_w) + H(b_0 || 0 || ... || 0)
	//  = H(b_0 + b_0 || b_1 + 0 || ... || b_w + 0)
	//  = H(    0     || b_1 || ...     || b_w)
	//
	// Afterwards a new byte can be shifted in.
	for b := 0; b < 256; b++ {
		var h Pol

		h = appendByte(h, byte(b), c.pol)
		for i := 0; i < windowSize-1; i++ {
			h = appendByte(h, 0, c.pol)
		}
		c.tables.out[b] = h
	}

	// calculate table for reduction mod Polynomial
	k := c.pol.Deg()
	for b := 0; b < 256; b++ {
		// mod_table[b] = A | B, where A = (b(x) * x^k mod pol) and  B = b(x) * x^k
		//
		// The 8 bits above deg(Polynomial) determine what happens next and so
		// these bits are used as a lookup to this table. The value is split in
		// two parts: Part A contains the result of the modulus operation, part
		// B is used to cancel out the 8 top bits so that one XOR operation is
		// enough to reduce modulo Polynomial
		c.tables.mod[b] = Pol(uint64(b)<<uint(k)).Mod(c.pol) | (Pol(b) << uint(k))
	}

	cache.entries[c.pol] = c.tables
}

// Next returns the position and length of the next chunk of data. If an error
// occurs while reading, the error is returned. Afterwards, the state of the
// current chunk is undefined. When the last chunk has been returned, all
// subsequent calls yield an io.EOF error.
func (c *Chunker) Next(data []byte) (Chunk, error) {
	data = data[:0]
	if !c.tablesInitialized {
		return Chunk{}, errors.New("tables for polynomial computation not initialized")
	}

	tabout := c.tables.out
	tabmod := c.tables.mod
	polShift := c.polShift
	minSize := c.MinSize
	maxSize := c.MaxSize
	buf := c.buf
	for {
		if c.bpos >= c.bmax {
			n, err := io.ReadFull(c.rd, buf[:])

			if err == io.ErrUnexpectedEOF {
				err = nil
			}

			// io.ReadFull only returns io.EOF when no bytes could be read. If
			// this is the case and we're in this branch, there are no more
			// bytes to buffer, so this was the last chunk. If a different
			// error has occurred, return that error and abandon the current
			// chunk.
			if err == io.EOF && !c.closed {
				c.closed = true

				// return current chunk, if any bytes have been processed
				if c.count > 0 {
					return Chunk{
						Start:  c.start,
						Length: c.count,
						Cut:    c.digest,
						Data:   data,
					}, nil
				}
			}

			if err != nil {
				return Chunk{}, err
			}

			c.bpos = 0
			c.bmax = uint(n)
		}

		// check if bytes have to be dismissed before starting a new chunk
		if c.pre > 0 {
			n := c.bmax - c.bpos
			if c.pre > uint(n) {
				c.pre -= uint(n)
				data = append(data, buf[c.bpos:c.bmax]...)

				c.count += uint(n)
				c.pos += uint(n)
				c.bpos = c.bmax

				continue
			}

			data = append(data, buf[c.bpos:c.bpos+c.pre]...)

			c.bpos += c.pre
			c.count += c.pre
			c.pos += c.pre
			c.pre = 0
		}

		add := c.count
		digest := c.digest
		win := c.window
		wpos := c.wpos
		for _, b := range buf[c.bpos:c.bmax] {
			// slide(b)
			out := win[wpos]
			win[wpos] = b
			digest ^= uint64(tabout[out])
			wpos++
			if wpos >= windowSize {
				wpos = 0
			}

			// updateDigest
			index := byte(digest >> polShift)
			digest <<= 8
			digest |= uint64(b)

			digest ^= uint64(tabmod[index])
			// end manual inline

			add++
			if add < minSize {
				continue
			}

			if (digest&c.splitmask) == 0 || add >= maxSize {
				i := add - c.count - 1
				data = append(data, c.buf[c.bpos:c.bpos+uint(i)+1]...)
				c.count = add
				c.pos += uint(i) + 1
				c.bpos += uint(i) + 1
				c.buf = buf

				chunk := Chunk{
					Start:  c.start,
					Length: c.count,
					Cut:    digest,
					Data:   data,
				}

				c.reset()

				return chunk, nil
			}
		}
		c.digest = digest
		c.window = win
		c.wpos = wpos

		steps := c.bmax - c.bpos
		if steps > 0 {
			data = append(data, c.buf[c.bpos:c.bpos+steps]...)
		}
		c.count += steps
		c.pos += steps
		c.bpos = c.bmax
	}
}

func updateDigest(digest uint64, polShift uint, tab tables, b byte) (newDigest uint64) {
	index := digest >> polShift
	digest <<= 8
	digest |= uint64(b)

	digest ^= uint64(tab.mod[index])
	return digest
}

func (c *Chunker) slide(digest uint64, b byte) (newDigest uint64) {
	out := c.window[c.wpos]
	c.window[c.wpos] = b
	digest ^= uint64(c.tables.out[out])
	c.wpos = (c.wpos + 1) % windowSize

	digest = updateDigest(digest, c.polShift, c.tables, b)
	return digest
}

func appendByte(hash Pol, b byte, pol Pol) Pol {
	hash <<= 8
	hash |= Pol(b)

	return hash.Mod(pol)
}
<commit_msg>Use uint to track window position<commit_after>package chunker

import (
	"errors"
	"io"
	"sync"
)

const (
	kiB = 1024
	miB = 1024 * kiB

	// WindowSize is the size of the sliding window.
	windowSize = 64

	// MinSize is the default minimal size of a chunk.
	MinSize = 512 * kiB
	// MaxSize is the default maximal size of a chunk.
	MaxSize = 8 * miB

	chunkerBufSize = 512 * kiB
)

type tables struct {
	out [256]Pol
	mod [256]Pol
}

// cache precomputed tables, these are read-only anyway
var cache struct {
	entries map[Pol]tables
	sync.Mutex
}

func init() {
	cache.entries = make(map[Pol]tables)
}

// Chunk is one content-dependent chunk of bytes whose end was cut when the
// Rabin Fingerprint had the value stored in Cut.
type Chunk struct {
	Start  uint
	Length uint
	Cut    uint64
	Data   []byte
}

type chunkerState struct {
	window [windowSize]byte
	wpos   uint

	buf  []byte
	bpos uint
	bmax uint

	start uint
	count uint
	pos   uint

	pre uint // wait for this many bytes before start calculating an new chunk

	digest uint64
}

type chunkerConfig struct {
	MinSize, MaxSize uint

	pol               Pol
	polShift          uint
	tables            tables
	tablesInitialized bool
	splitmask         uint64

	rd     io.Reader
	closed bool
}

// Chunker splits content with Rabin Fingerprints.
type Chunker struct {
	chunkerConfig
	chunkerState
}

// SetAverageBits allows to control the frequency of chunk discovery:
// the lower averageBits, the higher amount of chunks will be identified.
// The default value is 20 bits, so chunks will be of 1MiB size on average.
func (c *Chunker) SetAverageBits(averageBits int) {
	c.splitmask = (1 << uint64(averageBits)) - 1
}

// New returns a new Chunker based on polynomial p that reads from rd.
func New(rd io.Reader, pol Pol) *Chunker {
	return NewWithBoundaries(rd, pol, MinSize, MaxSize)
}

// NewWithBoundaries returns a new Chunker based on polynomial p that reads from
// rd and custom min and max size boundaries.
func NewWithBoundaries(rd io.Reader, pol Pol, min, max uint) *Chunker {
	c := &Chunker{
		chunkerState: chunkerState{
			buf: make([]byte, chunkerBufSize),
		},
		chunkerConfig: chunkerConfig{
			pol:       pol,
			rd:        rd,
			MinSize:   min,
			MaxSize:   max,
			splitmask: (1 << 20) - 1, // aim to create chunks of 20 bits or about 1MiB on average.
		},
	}

	c.reset()

	return c
}

// Reset reinitializes the chunker with a new reader and polynomial.
func (c *Chunker) Reset(rd io.Reader, pol Pol) {
	c.ResetWithBoundaries(rd, pol, MinSize, MaxSize)
}

// ResetWithBoundaries reinitializes the chunker with a new reader, polynomial
// and custom min and max size boundaries.
func (c *Chunker) ResetWithBoundaries(rd io.Reader, pol Pol, min, max uint) {
	*c = Chunker{
		chunkerState: chunkerState{
			buf: c.buf,
		},
		chunkerConfig: chunkerConfig{
			pol:       pol,
			rd:        rd,
			MinSize:   min,
			MaxSize:   max,
			splitmask: (1 << 20) - 1,
		},
	}

	c.reset()
}

func (c *Chunker) reset() {
	c.polShift = uint(c.pol.Deg() - 8)
	c.fillTables()

	for i := 0; i < windowSize; i++ {
		c.window[i] = 0
	}

	c.closed = false
	c.digest = 0
	c.wpos = 0
	c.count = 0
	c.digest = c.slide(c.digest, 1)
	c.start = c.pos

	// do not start a new chunk unless at least MinSize bytes have been read
	c.pre = c.MinSize - windowSize
}

// fillTables calculates out_table and mod_table for optimization. This
// implementation uses a cache in the global variable cache.
func (c *Chunker) fillTables() {
	// if polynomial hasn't been specified, do not compute anything for now
	if c.pol == 0 {
		return
	}

	c.tablesInitialized = true

	// test if the tables are cached for this polynomial
	cache.Lock()
	defer cache.Unlock()
	if t, ok := cache.entries[c.pol]; ok {
		c.tables = t
		return
	}

	// calculate table for sliding out bytes. The byte to slide out is used as
	// the index for the table, the value contains the following:
	// out_table[b] = Hash(b || 0 ||        ...        || 0)
	//                          \ windowsize-1 zero bytes /
	// To slide out byte b_0 for window size w with known hash
	// H := H(b_0 || ... || b_w), it is sufficient to add out_table[b_0]:
	//    H(b_0 || ... || b_w) + H(b_0 || 0 || ... || 0)
	//  = H(b_0 + b_0 || b_1 + 0 || ... || b_w + 0)
	//  = H(    0     || b_1 || ...     || b_w)
	//
	// Afterwards a new byte can be shifted in.
	for b := 0; b < 256; b++ {
		var h Pol

		h = appendByte(h, byte(b), c.pol)
		for i := 0; i < windowSize-1; i++ {
			h = appendByte(h, 0, c.pol)
		}
		c.tables.out[b] = h
	}

	// calculate table for reduction mod Polynomial
	k := c.pol.Deg()
	for b := 0; b < 256; b++ {
		// mod_table[b] = A | B, where A = (b(x) * x^k mod pol) and  B = b(x) * x^k
		//
		// The 8 bits above deg(Polynomial) determine what happens next and so
		// these bits are used as a lookup to this table. The value is split in
		// two parts: Part A contains the result of the modulus operation, part
		// B is used to cancel out the 8 top bits so that one XOR operation is
		// enough to reduce modulo Polynomial
		c.tables.mod[b] = Pol(uint64(b)<<uint(k)).Mod(c.pol) | (Pol(b) << uint(k))
	}

	cache.entries[c.pol] = c.tables
}

// Next returns the position and length of the next chunk of data. If an error
// occurs while reading, the error is returned. Afterwards, the state of the
// current chunk is undefined. When the last chunk has been returned, all
// subsequent calls yield an io.EOF error.
func (c *Chunker) Next(data []byte) (Chunk, error) {
	data = data[:0]
	if !c.tablesInitialized {
		return Chunk{}, errors.New("tables for polynomial computation not initialized")
	}

	tabout := c.tables.out
	tabmod := c.tables.mod
	polShift := c.polShift
	minSize := c.MinSize
	maxSize := c.MaxSize
	buf := c.buf
	for {
		if c.bpos >= c.bmax {
			n, err := io.ReadFull(c.rd, buf[:])

			if err == io.ErrUnexpectedEOF {
				err = nil
			}

			// io.ReadFull only returns io.EOF when no bytes could be read. If
			// this is the case and we're in this branch, there are no more
			// bytes to buffer, so this was the last chunk. If a different
			// error has occurred, return that error and abandon the current
			// chunk.
			if err == io.EOF && !c.closed {
				c.closed = true

				// return current chunk, if any bytes have been processed
				if c.count > 0 {
					return Chunk{
						Start:  c.start,
						Length: c.count,
						Cut:    c.digest,
						Data:   data,
					}, nil
				}
			}

			if err != nil {
				return Chunk{}, err
			}

			c.bpos = 0
			c.bmax = uint(n)
		}

		// check if bytes have to be dismissed before starting a new chunk
		if c.pre > 0 {
			n := c.bmax - c.bpos
			if c.pre > uint(n) {
				c.pre -= uint(n)
				data = append(data, buf[c.bpos:c.bmax]...)

				c.count += uint(n)
				c.pos += uint(n)
				c.bpos = c.bmax

				continue
			}

			data = append(data, buf[c.bpos:c.bpos+c.pre]...)

			c.bpos += c.pre
			c.count += c.pre
			c.pos += c.pre
			c.pre = 0
		}

		add := c.count
		digest := c.digest
		win := c.window
		wpos := c.wpos
		for _, b := range buf[c.bpos:c.bmax] {
			// slide(b)
			out := win[wpos]
			win[wpos] = b
			digest ^= uint64(tabout[out])
			wpos++
			wpos = wpos % windowSize

			// updateDigest
			index := byte(digest >> polShift)
			digest <<= 8
			digest |= uint64(b)

			digest ^= uint64(tabmod[index])
			// end manual inline

			add++
			if add < minSize {
				continue
			}

			if (digest&c.splitmask) == 0 || add >= maxSize {
				i := add - c.count - 1
				data = append(data, c.buf[c.bpos:c.bpos+uint(i)+1]...)
				c.count = add
				c.pos += uint(i) + 1
				c.bpos += uint(i) + 1
				c.buf = buf

				chunk := Chunk{
					Start:  c.start,
					Length: c.count,
					Cut:    digest,
					Data:   data,
				}

				c.reset()

				return chunk, nil
			}
		}
		c.digest = digest
		c.window = win
		c.wpos = wpos

		steps := c.bmax - c.bpos
		if steps > 0 {
			data = append(data, c.buf[c.bpos:c.bpos+steps]...)
		}
		c.count += steps
		c.pos += steps
		c.bpos = c.bmax
	}
}

func updateDigest(digest uint64, polShift uint, tab tables, b byte) (newDigest uint64) {
	index := digest >> polShift
	digest <<= 8
	digest |= uint64(b)

	digest ^= uint64(tab.mod[index])
	return digest
}

func (c *Chunker) slide(digest uint64, b byte) (newDigest uint64) {
	out := c.window[c.wpos]
	c.window[c.wpos] = b
	digest ^= uint64(c.tables.out[out])
	c.wpos = (c.wpos + 1) % windowSize

	digest = updateDigest(digest, c.polShift, c.tables, b)
	return digest
}

func appendByte(hash Pol, b byte, pol Pol) Pol {
	hash <<= 8
	hash |= Pol(b)

	return hash.Mod(pol)
}
<|endoftext|>


<commit_before>package router

import (
	"github.com/hashicorp/consul/agent/metadata"
	"github.com/hashicorp/consul/types"
	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/serf/serf"
)

// routerFn selects one of the router operations to map to incoming Serf events.
type routerFn func(types.AreaID, *metadata.Server) error

// handleMemberEvents attempts to apply the given Serf member event to the given
// router function.
func handleMemberEvent(logger hclog.Logger, fn routerFn, areaID types.AreaID, e serf.Event) {
	me, ok := e.(serf.MemberEvent)
	if !ok {
		logger.Error("Bad event type", "event", e)
		return
	}

	for _, m := range me.Members {
		ok, parts := metadata.IsConsulServer(m)
		if !ok {
			logger.Warn("Non-server in server-only area",
				"non_server", m.Name,
				"area", areaID,
			)
			continue
		}

		if err := fn(areaID, parts); err != nil {
			logger.Error("Failed to process event for server in area",
				"event", me.Type.String(),
				"server", m.Name,
				"area", areaID,
				"error", err,
			)
			continue
		}

		logger.Info("Handled event for server in area",
			"event", me.Type.String(),
			"server", m.Name,
			"area", areaID,
		)
	}
}

// HandleSerfEvents is a long-running goroutine that pushes incoming events from
// a Serf manager's channel into the given router. This will return when the
// shutdown channel is closed.
func HandleSerfEvents(logger hclog.Logger, router *Router, areaID types.AreaID, shutdownCh <-chan struct{}, eventCh <-chan serf.Event) {
	for {
		select {
		case <-shutdownCh:
			return

		case e := <-eventCh:
			switch e.EventType() {
			case serf.EventMemberJoin:
				handleMemberEvent(logger, router.AddServer, areaID, e)

			case serf.EventMemberLeave, serf.EventMemberReap:
				handleMemberEvent(logger, router.RemoveServer, areaID, e)

			case serf.EventMemberFailed:
				handleMemberEvent(logger, router.FailServer, areaID, e)

			// All of these event types are ignored.
			case serf.EventMemberUpdate:
				handleMemberEvent(logger, router.AddServer, areaID, e)
			case serf.EventUser:
			case serf.EventQuery:

			default:
				logger.Warn("Unhandled Serf Event", "event", e)
			}
		}
	}
}
<commit_msg>Fixed comment on wrong line.<commit_after>package router

import (
	"github.com/hashicorp/consul/agent/metadata"
	"github.com/hashicorp/consul/types"
	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/serf/serf"
)

// routerFn selects one of the router operations to map to incoming Serf events.
type routerFn func(types.AreaID, *metadata.Server) error

// handleMemberEvents attempts to apply the given Serf member event to the given
// router function.
func handleMemberEvent(logger hclog.Logger, fn routerFn, areaID types.AreaID, e serf.Event) {
	me, ok := e.(serf.MemberEvent)
	if !ok {
		logger.Error("Bad event type", "event", e)
		return
	}

	for _, m := range me.Members {
		ok, parts := metadata.IsConsulServer(m)
		if !ok {
			logger.Warn("Non-server in server-only area",
				"non_server", m.Name,
				"area", areaID,
			)
			continue
		}

		if err := fn(areaID, parts); err != nil {
			logger.Error("Failed to process event for server in area",
				"event", me.Type.String(),
				"server", m.Name,
				"area", areaID,
				"error", err,
			)
			continue
		}

		logger.Info("Handled event for server in area",
			"event", me.Type.String(),
			"server", m.Name,
			"area", areaID,
		)
	}
}

// HandleSerfEvents is a long-running goroutine that pushes incoming events from
// a Serf manager's channel into the given router. This will return when the
// shutdown channel is closed.
func HandleSerfEvents(logger hclog.Logger, router *Router, areaID types.AreaID, shutdownCh <-chan struct{}, eventCh <-chan serf.Event) {
	for {
		select {
		case <-shutdownCh:
			return

		case e := <-eventCh:
			switch e.EventType() {
			case serf.EventMemberJoin:
				handleMemberEvent(logger, router.AddServer, areaID, e)

			case serf.EventMemberLeave, serf.EventMemberReap:
				handleMemberEvent(logger, router.RemoveServer, areaID, e)

			case serf.EventMemberFailed:
				handleMemberEvent(logger, router.FailServer, areaID, e)

			case serf.EventMemberUpdate:
				handleMemberEvent(logger, router.AddServer, areaID, e)

			// All of these event types are ignored.
			case serf.EventUser:
			case serf.EventQuery:

			default:
				logger.Warn("Unhandled Serf Event", "event", e)
			}
		}
	}
}
<|endoftext|>


<commit_before>package dev

import (
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
)

type transp interface {
	Read(b []byte) (n int, err error)
	Write(b []byte) (n int, err error)
	SetDeadline(t time.Time) error
	SetWriteDeadline(t time.Time) error
	Close() error
}

type transpTCP struct {
	net.Conn
}

type transpTelnet struct {
	net.Conn
	logger hasPrintf
}

type telnetOptions struct {
	supressGoAhead bool
	linemode       bool
}

func (s *transpTelnet) Read(b []byte) (int, error) {
	n1, err1 := s.Conn.Read(b)
	if err1 != nil {
		return n1, err1
	}
	n2, err2 := telnetNegotiation(b, n1, s, s.logger, false)
	return n2, err2
}

type transpPipe struct {
	proc     *exec.Cmd
	stdout   io.ReadCloser
	stderr   io.ReadCloser
	reader   io.Reader
	writer   io.WriteCloser
	logger   hasPrintf
	devLabel string
	debug    bool
	ctx      context.Context
	cancel   context.CancelFunc
}

func (s *transpPipe) result(n int, err error) (int, error) {
	var waitErr error
	if err != nil {
		waitErr = s.proc.Wait()
	}

	ctxErr := s.ctx.Err()

	if ctxErr != nil || waitErr != nil {
		return n, fmt.Errorf("transPipe result error: error=[%v] context=[%v] wait=[%v]", err, ctxErr, waitErr)
	}

	return n, err
}

func (s *transpPipe) Read(b []byte) (int, error) {
	n, err := s.reader.Read(b)
	return s.result(n, err)
}

func (s *transpPipe) Write(b []byte) (int, error) {
	n, err := s.writer.Write(b)
	return s.result(n, err)
}

func (s *transpPipe) SetDeadline(t time.Time) error {
	s.logger.Printf("transpPipe.SetDeadline: FIXME WRITEME")
	return nil
}

func (s *transpPipe) SetWriteDeadline(t time.Time) error {
	s.logger.Printf("transpPipe.SetWriteDeadline: FIXME WRITEME")
	return nil
}

func (s *transpPipe) Close() error {

	if s.debug {
		s.logger.Printf("transpPipe.Close: %s contextErr=[%v]", s.devLabel, s.ctx.Err())
	}

	s.cancel()

	err1 := s.stdout.Close()
	err2 := s.stderr.Close()
	err3 := s.writer.Close()

	if err1 != nil || err2 != nil || err3 != nil {
		return fmt.Errorf("transpPipe: close error: out=[%v] err=[%v] writer=[%v]", err1, err2, err3)
	}

	return nil
}

type transpSSH struct {
	devLabel string
	conn     net.Conn
	client   *ssh.Client
	session  *ssh.Session
	writer   io.Writer
	reader   io.Reader
}

func (s *transpSSH) Read(b []byte) (int, error) {
	return s.reader.Read(b)
}

func (s *transpSSH) Write(b []byte) (int, error) {

	n, err := s.writer.Write(b)
	if err != nil {
		return -1, fmt.Errorf("ssh write(%s): %v", b, err)
	}

	return n, nil
}

func (s *transpSSH) SetDeadline(t time.Time) error {
	return s.conn.SetDeadline(t)
}

func (s *transpSSH) SetWriteDeadline(t time.Time) error {
	return s.conn.SetWriteDeadline(t)
}

func (s *transpSSH) Close() error {
	err1 := s.session.Close()
	err2 := s.conn.Close()
	if err1 != nil || err2 != nil {
		return fmt.Errorf("close error: session=[%v] conn=[%v]", err1, err2)
	}
	return nil
}

func openTransportPipe(logger hasPrintf, modelName, devID, hostPort, transports, user, pass string, args []string, debug bool, timeout time.Duration) (transp, string, bool, error) {
	s, err := openPipe(logger, modelName, devID, hostPort, transports, user, pass, args, debug, timeout)
	return s, "pipe", true, err
}

func openPipe(logger hasPrintf, modelName, devID, hostPort, transports, user, pass string, args []string, debug bool, timeout time.Duration) (transp, error) {

	devLabel := fmt.Sprintf("%s %s %s", modelName, devID, hostPort)

	logger.Printf("openPipe: %s - opening", devLabel)

	ctx, cancel := context.WithTimeout(context.Background(), timeout)

	c := exec.CommandContext(ctx, args[0], args[1:]...)

	pipeOut, outErr := c.StdoutPipe()
	if outErr != nil {
		cancel()
		return nil, fmt.Errorf("openPipe: StdoutPipe: %s - %v", devLabel, outErr)
	}

	pipeErr, errErr := c.StderrPipe()
	if errErr != nil {
		cancel()
		return nil, fmt.Errorf("openPipe: StderrPipe: %s - %v", devLabel, errErr)
	}

	writer, wrErr := c.StdinPipe()
	if wrErr != nil {
		cancel()
		return nil, fmt.Errorf("openPipe: StdinPipe: %s - %v", devLabel, wrErr)
	}

	s := &transpPipe{proc: c, logger: logger, devLabel: devLabel, debug: debug}
	s.reader = io.MultiReader(pipeOut, pipeErr)
	s.stdout = pipeOut
	s.stderr = pipeErr
	s.writer = writer
	s.ctx = ctx
	s.cancel = cancel

	logger.Printf("openPipe: %s - starting", devLabel)

	os.Setenv("JAZIGO_DEV_ID", devID)
	os.Setenv("JAZIGO_DEV_HOSTPORT", hostPort)
	os.Setenv("JAZIGO_DEV_USER", user)
	os.Setenv("JAZIGO_DEV_PASS", pass)

	if startErr := s.proc.Start(); startErr != nil {
		s.Close()
		return nil, fmt.Errorf("openPipe: error: %v", startErr)
	}

	logger.Printf("openPipe: %s - started", devLabel)

	return s, nil
}

func openTransport(logger hasPrintf, modelName, devID, hostPort, transports, user, pass string) (transp, string, bool, error) {
	tList := strings.Split(transports, ",")
	if len(tList) < 1 {
		return nil, transports, false, fmt.Errorf("openTransport: missing transports: [%s]", transports)
	}

	var lastErr error

	timeout := 10 * time.Second

	for _, t := range tList {
		switch t {
		case "ssh":
			hp := forceHostPort(hostPort, "22")
			s, err := openSSH(logger, modelName, devID, hp, timeout, user, pass)
			if err == nil {
				return s, t, true, nil
			}
			logger.Printf("openTransport: %v", err)
			lastErr = err
		case "telnet":
			hp := forceHostPort(hostPort, "23")
			s, err := openTelnet(logger, modelName, devID, hp, timeout)
			if err == nil {
				return s, t, false, nil
			}
			logger.Printf("openTransport: %v", err)
			lastErr = err
		default:
			s, err := openTCP(logger, modelName, devID, hostPort, timeout)
			if err == nil {
				return s, t, false, nil
			}
			logger.Printf("openTransport: %v", err)
			lastErr = err
		}
	}

	return nil, transports, false, fmt.Errorf("openTransport: %s %s %s %s - unable to open transport: last error: %v", modelName, devID, hostPort, transports, lastErr)
}

func forceHostPort(hostPort, defaultPort string) string {
	i := strings.Index(hostPort, ":")
	if i < 0 {
		return fmt.Sprintf("%s:%s", hostPort, defaultPort)
	}
	return hostPort
}

func openSSH(logger hasPrintf, modelName, devID, hostPort string, timeout time.Duration, user, pass string) (transp, error) {

	conn, dialErr := net.DialTimeout("tcp", hostPort, timeout)
	if dialErr != nil {
		return nil, fmt.Errorf("openSSH: Dial: %s %s %s - %v", modelName, devID, hostPort, dialErr)
	}

	conf := &ssh.Config{}
	conf.SetDefaults()
	conf.Ciphers = append(conf.Ciphers, "3des-cbc") // 3des-cbc is needed for IOS XR

	config := &ssh.ClientConfig{
		Config: *conf,
		User:   user,
		Auth: []ssh.AuthMethod{
			ssh.Password(pass),
		},
		Timeout: timeout,
	}

	c, chans, reqs, connErr := ssh.NewClientConn(conn, hostPort, config)
	if connErr != nil {
		return nil, fmt.Errorf("openSSH: NewClientConn: %s %s %s - %v", modelName, devID, hostPort, connErr)
	}

	cli := ssh.NewClient(c, chans, reqs)

	s := &transpSSH{conn: conn, client: cli, devLabel: fmt.Sprintf("%s %s %s", modelName, devID, hostPort) /*, logger: logger*/}

	ses, sessionErr := s.client.NewSession()
	if sessionErr != nil {
		return nil, fmt.Errorf("openSSH: NewSession: %s - %v", s.devLabel, sessionErr)
	}

	s.session = ses

	modes := ssh.TerminalModes{
		ssh.ECHO: 0, // disable echoing
	}

	if ptyErr := ses.RequestPty("xterm", 80, 40, modes); ptyErr != nil {
		return nil, fmt.Errorf("openSSH: Pty: %s - %v", s.devLabel, ptyErr)
	}

	pipeOut, outErr := ses.StdoutPipe()
	if outErr != nil {
		return nil, fmt.Errorf("openSSH: StdoutPipe: %s - %v", s.devLabel, outErr)
	}

	pipeErr, errErr := ses.StderrPipe()
	if errErr != nil {
		return nil, fmt.Errorf("openSSH: StderrPipe: %s - %v", s.devLabel, errErr)
	}

	s.reader = io.MultiReader(pipeOut, pipeErr)

	writer, wrErr := ses.StdinPipe()
	if wrErr != nil {
		return nil, fmt.Errorf("openSSH: StdinPipe: %s - %v", s.devLabel, wrErr)
	}

	s.writer = writer

	if shellErr := ses.Shell(); shellErr != nil {
		return nil, fmt.Errorf("openSSH: Remote shell error: %s - %v", s.devLabel, shellErr)
	}

	return s, nil
}

func openTelnet(logger hasPrintf, modelName, devID, hostPort string, timeout time.Duration) (transp, error) {

	conn, err := net.DialTimeout("tcp", hostPort, timeout)
	if err != nil {
		return nil, fmt.Errorf("openTelnet: %s %s %s - %v", modelName, devID, hostPort, err)
	}

	return &transpTelnet{conn, logger}, nil
}

func openTCP(logger hasPrintf, modelName, devID, hostPort string, timeout time.Duration) (transp, error) {

	conn, err := net.DialTimeout("tcp", hostPort, timeout)
	if err != nil {
		return nil, fmt.Errorf("openTCP: %s %s %s - %v", modelName, devID, hostPort, err)
	}

	return &transpTCP{conn}, nil
}
<commit_msg>Provide ClientConfig.HostKeyCallback<commit_after>package dev

import (
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
)

type transp interface {
	Read(b []byte) (n int, err error)
	Write(b []byte) (n int, err error)
	SetDeadline(t time.Time) error
	SetWriteDeadline(t time.Time) error
	Close() error
}

type transpTCP struct {
	net.Conn
}

type transpTelnet struct {
	net.Conn
	logger hasPrintf
}

type telnetOptions struct {
	supressGoAhead bool
	linemode       bool
}

func (s *transpTelnet) Read(b []byte) (int, error) {
	n1, err1 := s.Conn.Read(b)
	if err1 != nil {
		return n1, err1
	}
	n2, err2 := telnetNegotiation(b, n1, s, s.logger, false)
	return n2, err2
}

type transpPipe struct {
	proc     *exec.Cmd
	stdout   io.ReadCloser
	stderr   io.ReadCloser
	reader   io.Reader
	writer   io.WriteCloser
	logger   hasPrintf
	devLabel string
	debug    bool
	ctx      context.Context
	cancel   context.CancelFunc
}

func (s *transpPipe) result(n int, err error) (int, error) {
	var waitErr error
	if err != nil {
		waitErr = s.proc.Wait()
	}

	ctxErr := s.ctx.Err()

	if ctxErr != nil || waitErr != nil {
		return n, fmt.Errorf("transPipe result error: error=[%v] context=[%v] wait=[%v]", err, ctxErr, waitErr)
	}

	return n, err
}

func (s *transpPipe) Read(b []byte) (int, error) {
	n, err := s.reader.Read(b)
	return s.result(n, err)
}

func (s *transpPipe) Write(b []byte) (int, error) {
	n, err := s.writer.Write(b)
	return s.result(n, err)
}

func (s *transpPipe) SetDeadline(t time.Time) error {
	s.logger.Printf("transpPipe.SetDeadline: FIXME WRITEME")
	return nil
}

func (s *transpPipe) SetWriteDeadline(t time.Time) error {
	s.logger.Printf("transpPipe.SetWriteDeadline: FIXME WRITEME")
	return nil
}

func (s *transpPipe) Close() error {

	if s.debug {
		s.logger.Printf("transpPipe.Close: %s contextErr=[%v]", s.devLabel, s.ctx.Err())
	}

	s.cancel()

	err1 := s.stdout.Close()
	err2 := s.stderr.Close()
	err3 := s.writer.Close()

	if err1 != nil || err2 != nil || err3 != nil {
		return fmt.Errorf("transpPipe: close error: out=[%v] err=[%v] writer=[%v]", err1, err2, err3)
	}

	return nil
}

type transpSSH struct {
	devLabel string
	conn     net.Conn
	client   *ssh.Client
	session  *ssh.Session
	writer   io.Writer
	reader   io.Reader
}

func (s *transpSSH) Read(b []byte) (int, error) {
	return s.reader.Read(b)
}

func (s *transpSSH) Write(b []byte) (int, error) {

	n, err := s.writer.Write(b)
	if err != nil {
		return -1, fmt.Errorf("ssh write(%s): %v", b, err)
	}

	return n, nil
}

func (s *transpSSH) SetDeadline(t time.Time) error {
	return s.conn.SetDeadline(t)
}

func (s *transpSSH) SetWriteDeadline(t time.Time) error {
	return s.conn.SetWriteDeadline(t)
}

func (s *transpSSH) Close() error {
	err1 := s.session.Close()
	err2 := s.conn.Close()
	if err1 != nil || err2 != nil {
		return fmt.Errorf("close error: session=[%v] conn=[%v]", err1, err2)
	}
	return nil
}

func openTransportPipe(logger hasPrintf, modelName, devID, hostPort, transports, user, pass string, args []string, debug bool, timeout time.Duration) (transp, string, bool, error) {
	s, err := openPipe(logger, modelName, devID, hostPort, transports, user, pass, args, debug, timeout)
	return s, "pipe", true, err
}

func openPipe(logger hasPrintf, modelName, devID, hostPort, transports, user, pass string, args []string, debug bool, timeout time.Duration) (transp, error) {

	devLabel := fmt.Sprintf("%s %s %s", modelName, devID, hostPort)

	logger.Printf("openPipe: %s - opening", devLabel)

	ctx, cancel := context.WithTimeout(context.Background(), timeout)

	c := exec.CommandContext(ctx, args[0], args[1:]...)

	pipeOut, outErr := c.StdoutPipe()
	if outErr != nil {
		cancel()
		return nil, fmt.Errorf("openPipe: StdoutPipe: %s - %v", devLabel, outErr)
	}

	pipeErr, errErr := c.StderrPipe()
	if errErr != nil {
		cancel()
		return nil, fmt.Errorf("openPipe: StderrPipe: %s - %v", devLabel, errErr)
	}

	writer, wrErr := c.StdinPipe()
	if wrErr != nil {
		cancel()
		return nil, fmt.Errorf("openPipe: StdinPipe: %s - %v", devLabel, wrErr)
	}

	s := &transpPipe{proc: c, logger: logger, devLabel: devLabel, debug: debug}
	s.reader = io.MultiReader(pipeOut, pipeErr)
	s.stdout = pipeOut
	s.stderr = pipeErr
	s.writer = writer
	s.ctx = ctx
	s.cancel = cancel

	logger.Printf("openPipe: %s - starting", devLabel)

	os.Setenv("JAZIGO_DEV_ID", devID)
	os.Setenv("JAZIGO_DEV_HOSTPORT", hostPort)
	os.Setenv("JAZIGO_DEV_USER", user)
	os.Setenv("JAZIGO_DEV_PASS", pass)

	if startErr := s.proc.Start(); startErr != nil {
		s.Close()
		return nil, fmt.Errorf("openPipe: error: %v", startErr)
	}

	logger.Printf("openPipe: %s - started", devLabel)

	return s, nil
}

func openTransport(logger hasPrintf, modelName, devID, hostPort, transports, user, pass string) (transp, string, bool, error) {
	tList := strings.Split(transports, ",")
	if len(tList) < 1 {
		return nil, transports, false, fmt.Errorf("openTransport: missing transports: [%s]", transports)
	}

	var lastErr error

	timeout := 10 * time.Second

	for _, t := range tList {
		switch t {
		case "ssh":
			hp := forceHostPort(hostPort, "22")
			s, err := openSSH(logger, modelName, devID, hp, timeout, user, pass)
			if err == nil {
				return s, t, true, nil
			}
			logger.Printf("openTransport: %v", err)
			lastErr = err
		case "telnet":
			hp := forceHostPort(hostPort, "23")
			s, err := openTelnet(logger, modelName, devID, hp, timeout)
			if err == nil {
				return s, t, false, nil
			}
			logger.Printf("openTransport: %v", err)
			lastErr = err
		default:
			s, err := openTCP(logger, modelName, devID, hostPort, timeout)
			if err == nil {
				return s, t, false, nil
			}
			logger.Printf("openTransport: %v", err)
			lastErr = err
		}
	}

	return nil, transports, false, fmt.Errorf("openTransport: %s %s %s %s - unable to open transport: last error: %v", modelName, devID, hostPort, transports, lastErr)
}

func forceHostPort(hostPort, defaultPort string) string {
	i := strings.Index(hostPort, ":")
	if i < 0 {
		return fmt.Sprintf("%s:%s", hostPort, defaultPort)
	}
	return hostPort
}

func hostKeyCheck(hostname string, remote net.Addr, key ssh.PublicKey) error {
	return nil // FIXME hostKeyCheck accept anything
}

func openSSH(logger hasPrintf, modelName, devID, hostPort string, timeout time.Duration, user, pass string) (transp, error) {

	conn, dialErr := net.DialTimeout("tcp", hostPort, timeout)
	if dialErr != nil {
		return nil, fmt.Errorf("openSSH: Dial: %s %s %s - %v", modelName, devID, hostPort, dialErr)
	}

	conf := &ssh.Config{}
	conf.SetDefaults()
	conf.Ciphers = append(conf.Ciphers, "3des-cbc") // 3des-cbc is needed for IOS XR

	config := &ssh.ClientConfig{
		Config: *conf,
		User:   user,
		Auth: []ssh.AuthMethod{
			ssh.Password(pass),
		},
		Timeout:         timeout,
		HostKeyCallback: hostKeyCheck,
	}

	c, chans, reqs, connErr := ssh.NewClientConn(conn, hostPort, config)
	if connErr != nil {
		return nil, fmt.Errorf("openSSH: NewClientConn: %s %s %s - %v", modelName, devID, hostPort, connErr)
	}

	cli := ssh.NewClient(c, chans, reqs)

	s := &transpSSH{conn: conn, client: cli, devLabel: fmt.Sprintf("%s %s %s", modelName, devID, hostPort) /*, logger: logger*/}

	ses, sessionErr := s.client.NewSession()
	if sessionErr != nil {
		return nil, fmt.Errorf("openSSH: NewSession: %s - %v", s.devLabel, sessionErr)
	}

	s.session = ses

	modes := ssh.TerminalModes{
		ssh.ECHO: 0, // disable echoing
	}

	if ptyErr := ses.RequestPty("xterm", 80, 40, modes); ptyErr != nil {
		return nil, fmt.Errorf("openSSH: Pty: %s - %v", s.devLabel, ptyErr)
	}

	pipeOut, outErr := ses.StdoutPipe()
	if outErr != nil {
		return nil, fmt.Errorf("openSSH: StdoutPipe: %s - %v", s.devLabel, outErr)
	}

	pipeErr, errErr := ses.StderrPipe()
	if errErr != nil {
		return nil, fmt.Errorf("openSSH: StderrPipe: %s - %v", s.devLabel, errErr)
	}

	s.reader = io.MultiReader(pipeOut, pipeErr)

	writer, wrErr := ses.StdinPipe()
	if wrErr != nil {
		return nil, fmt.Errorf("openSSH: StdinPipe: %s - %v", s.devLabel, wrErr)
	}

	s.writer = writer

	if shellErr := ses.Shell(); shellErr != nil {
		return nil, fmt.Errorf("openSSH: Remote shell error: %s - %v", s.devLabel, shellErr)
	}

	return s, nil
}

func openTelnet(logger hasPrintf, modelName, devID, hostPort string, timeout time.Duration) (transp, error) {

	conn, err := net.DialTimeout("tcp", hostPort, timeout)
	if err != nil {
		return nil, fmt.Errorf("openTelnet: %s %s %s - %v", modelName, devID, hostPort, err)
	}

	return &transpTelnet{conn, logger}, nil
}

func openTCP(logger hasPrintf, modelName, devID, hostPort string, timeout time.Duration) (transp, error) {

	conn, err := net.DialTimeout("tcp", hostPort, timeout)
	if err != nil {
		return nil, fmt.Errorf("openTCP: %s %s %s - %v", modelName, devID, hostPort, err)
	}

	return &transpTCP{conn}, nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package manual_test

import (
	"fmt"
	"io"
	"os"

	gc "launchpad.net/gocheck"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/bootstrap"
	"github.com/juju/juju/environs/cloudinit"
	"github.com/juju/juju/environs/filestorage"
	"github.com/juju/juju/environs/manual"
	"github.com/juju/juju/environs/storage"
	"github.com/juju/juju/environs/tools"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/juju/testing"
	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/version"
)

type bootstrapSuite struct {
	testing.JujuConnSuite
	env *localStorageEnviron
}

var _ = gc.Suite(&bootstrapSuite{})

type localStorageEnviron struct {
	environs.Environ
	storage     storage.Storage
	storageAddr string
	storageDir  string
}

func (e *localStorageEnviron) Storage() storage.Storage {
	return e.storage
}

func (e *localStorageEnviron) StorageAddr() string {
	return e.storageAddr
}

func (e *localStorageEnviron) StorageDir() string {
	return e.storageDir
}

func (s *bootstrapSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)
	s.env = &localStorageEnviron{
		Environ:    s.Conn.Environ,
		storageDir: c.MkDir(),
	}
	storage, err := filestorage.NewFileStorageWriter(s.env.storageDir)
	c.Assert(err, gc.IsNil)
	s.env.storage = storage
}

func (s *bootstrapSuite) getArgs(c *gc.C) manual.BootstrapArgs {
	hostname, err := os.Hostname()
	c.Assert(err, gc.IsNil)
	toolsList, err := tools.FindBootstrapTools(s.Conn.Environ, tools.BootstrapToolsParams{})
	c.Assert(err, gc.IsNil)
	arch := "amd64"
	return manual.BootstrapArgs{
		Host:          hostname,
		DataDir:       "/var/lib/juju",
		Environ:       s.env,
		PossibleTools: toolsList,
		Series:        coretesting.FakeDefaultSeries,
		HardwareCharacteristics: &instance.HardwareCharacteristics{
			Arch: &arch,
		},
		Context: coretesting.Context(c),
	}
}

func (s *bootstrapSuite) TestBootstrap(c *gc.C) {
	args := s.getArgs(c)
	args.Host = "ubuntu@" + args.Host

	defer fakeSSH{SkipDetection: true}.install(c).Restore()
	err := manual.Bootstrap(args)
	c.Assert(err, gc.IsNil)

	bootstrapState, err := bootstrap.LoadState(s.env.Storage())
	c.Assert(err, gc.IsNil)
	c.Assert(
		bootstrapState.StateInstances,
		gc.DeepEquals,
		[]instance.Id{manual.BootstrapInstanceId},
	)

	// Do it all again; this should work, despite the fact that
	// there's a bootstrap state file. Existence for that is
	// checked in general bootstrap code (environs/bootstrap).
	defer fakeSSH{SkipDetection: true}.install(c).Restore()
	err = manual.Bootstrap(args)
	c.Assert(err, gc.IsNil)

	// We *do* check that the machine has no juju* upstart jobs, though.
	defer fakeSSH{
		Provisioned:        true,
		SkipDetection:      true,
		SkipProvisionAgent: true,
	}.install(c).Restore()
	err = manual.Bootstrap(args)
	c.Assert(err, gc.Equals, manual.ErrProvisioned)
}

func (s *bootstrapSuite) TestBootstrapScriptFailure(c *gc.C) {
	args := s.getArgs(c)
	args.Host = "ubuntu@" + args.Host
	defer fakeSSH{SkipDetection: true, ProvisionAgentExitCode: 1}.install(c).Restore()
	err := manual.Bootstrap(args)
	c.Assert(err, gc.NotNil)

	// Since the script failed, the state file should have been
	// removed from storage.
	_, err = bootstrap.LoadState(s.env.Storage())
	c.Check(err, gc.Equals, environs.ErrNotBootstrapped)
}

func (s *bootstrapSuite) TestBootstrapEmptyDataDir(c *gc.C) {
	args := s.getArgs(c)
	args.DataDir = ""
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "data-dir argument is empty")
}

func (s *bootstrapSuite) TestBootstrapEmptyHost(c *gc.C) {
	args := s.getArgs(c)
	args.Host = ""
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "host argument is empty")
}

func (s *bootstrapSuite) TestBootstrapNilEnviron(c *gc.C) {
	args := s.getArgs(c)
	args.Environ = nil
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "environ argument is nil")
}

func (s *bootstrapSuite) TestBootstrapNoMatchingTools(c *gc.C) {
	// Empty tools list.
	args := s.getArgs(c)
	args.PossibleTools = nil
	defer fakeSSH{SkipDetection: true, SkipProvisionAgent: true}.install(c).Restore()
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "possible tools is empty")

	// Non-empty list, but none that match the series/arch.
	args = s.getArgs(c)
	args.Series = "edgy"
	defer fakeSSH{SkipDetection: true, SkipProvisionAgent: true}.install(c).Restore()
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "no matching tools available")
}

func (s *bootstrapSuite) TestBootstrapToolsFileURL(c *gc.C) {
	storageName := tools.StorageName(version.Current)
	sftpURL, err := s.env.Storage().URL(storageName)
	c.Assert(err, gc.IsNil)
	fileURL := fmt.Sprintf("file://%s/%s", s.env.storageDir, storageName)
	s.testBootstrapToolsURL(c, sftpURL, fileURL)
}

func (s *bootstrapSuite) TestBootstrapToolsExternalURL(c *gc.C) {
	const externalURL = "http://test.invalid/tools.tgz"
	s.testBootstrapToolsURL(c, externalURL, externalURL)
}

func (s *bootstrapSuite) testBootstrapToolsURL(c *gc.C, toolsURL, expectedURL string) {
	s.PatchValue(manual.ProvisionMachineAgent, func(host string, mcfg *cloudinit.MachineConfig, w io.Writer) error {
		c.Assert(mcfg.Tools.URL, gc.Equals, expectedURL)
		return nil
	})
	args := s.getArgs(c)
	args.PossibleTools[0].URL = toolsURL
	defer fakeSSH{SkipDetection: true}.install(c).Restore()
	err := manual.Bootstrap(args)
	c.Assert(err, gc.IsNil)
}
<commit_msg>Fix a manual test<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package manual_test

import (
	"fmt"
	"io"
	"os"

	gc "launchpad.net/gocheck"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/bootstrap"
	"github.com/juju/juju/environs/cloudinit"
	"github.com/juju/juju/environs/filestorage"
	"github.com/juju/juju/environs/manual"
	"github.com/juju/juju/environs/storage"
	"github.com/juju/juju/environs/tools"
	"github.com/juju/juju/instance"
	"github.com/juju/juju/juju/testing"
	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/version"
)

type bootstrapSuite struct {
	testing.JujuConnSuite
	env *localStorageEnviron
}

var _ = gc.Suite(&bootstrapSuite{})

type localStorageEnviron struct {
	environs.Environ
	storage     storage.Storage
	storageAddr string
	storageDir  string
}

func (e *localStorageEnviron) Storage() storage.Storage {
	return e.storage
}

func (e *localStorageEnviron) StorageAddr() string {
	return e.storageAddr
}

func (e *localStorageEnviron) StorageDir() string {
	return e.storageDir
}

func (s *bootstrapSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)
	s.env = &localStorageEnviron{
		Environ:    s.Conn.Environ,
		storageDir: c.MkDir(),
	}
	storage, err := filestorage.NewFileStorageWriter(s.env.storageDir)
	c.Assert(err, gc.IsNil)
	s.env.storage = storage
}

func (s *bootstrapSuite) getArgs(c *gc.C) manual.BootstrapArgs {
	hostname, err := os.Hostname()
	c.Assert(err, gc.IsNil)
	toolsList, err := tools.FindBootstrapTools(s.Conn.Environ, tools.BootstrapToolsParams{})
	c.Assert(err, gc.IsNil)
	arch := "amd64"
	return manual.BootstrapArgs{
		Host:          hostname,
		DataDir:       "/var/lib/juju",
		Environ:       s.env,
		PossibleTools: toolsList,
		Series:        toolsList[0].Version.Series,
		HardwareCharacteristics: &instance.HardwareCharacteristics{
			Arch: &arch,
		},
		Context: coretesting.Context(c),
	}
}

func (s *bootstrapSuite) TestBootstrap(c *gc.C) {
	args := s.getArgs(c)
	args.Host = "ubuntu@" + args.Host

	defer fakeSSH{SkipDetection: true}.install(c).Restore()
	err := manual.Bootstrap(args)
	c.Assert(err, gc.IsNil)

	bootstrapState, err := bootstrap.LoadState(s.env.Storage())
	c.Assert(err, gc.IsNil)
	c.Assert(
		bootstrapState.StateInstances,
		gc.DeepEquals,
		[]instance.Id{manual.BootstrapInstanceId},
	)

	// Do it all again; this should work, despite the fact that
	// there's a bootstrap state file. Existence for that is
	// checked in general bootstrap code (environs/bootstrap).
	defer fakeSSH{SkipDetection: true}.install(c).Restore()
	err = manual.Bootstrap(args)
	c.Assert(err, gc.IsNil)

	// We *do* check that the machine has no juju* upstart jobs, though.
	defer fakeSSH{
		Provisioned:        true,
		SkipDetection:      true,
		SkipProvisionAgent: true,
	}.install(c).Restore()
	err = manual.Bootstrap(args)
	c.Assert(err, gc.Equals, manual.ErrProvisioned)
}

func (s *bootstrapSuite) TestBootstrapScriptFailure(c *gc.C) {
	args := s.getArgs(c)
	args.Host = "ubuntu@" + args.Host
	defer fakeSSH{SkipDetection: true, ProvisionAgentExitCode: 1}.install(c).Restore()
	err := manual.Bootstrap(args)
	c.Assert(err, gc.NotNil)

	// Since the script failed, the state file should have been
	// removed from storage.
	_, err = bootstrap.LoadState(s.env.Storage())
	c.Check(err, gc.Equals, environs.ErrNotBootstrapped)
}

func (s *bootstrapSuite) TestBootstrapEmptyDataDir(c *gc.C) {
	args := s.getArgs(c)
	args.DataDir = ""
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "data-dir argument is empty")
}

func (s *bootstrapSuite) TestBootstrapEmptyHost(c *gc.C) {
	args := s.getArgs(c)
	args.Host = ""
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "host argument is empty")
}

func (s *bootstrapSuite) TestBootstrapNilEnviron(c *gc.C) {
	args := s.getArgs(c)
	args.Environ = nil
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "environ argument is nil")
}

func (s *bootstrapSuite) TestBootstrapNoMatchingTools(c *gc.C) {
	// Empty tools list.
	args := s.getArgs(c)
	args.PossibleTools = nil
	defer fakeSSH{SkipDetection: true, SkipProvisionAgent: true}.install(c).Restore()
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "possible tools is empty")

	// Non-empty list, but none that match the series/arch.
	args = s.getArgs(c)
	args.Series = "edgy"
	defer fakeSSH{SkipDetection: true, SkipProvisionAgent: true}.install(c).Restore()
	c.Assert(manual.Bootstrap(args), gc.ErrorMatches, "no matching tools available")
}

func (s *bootstrapSuite) TestBootstrapToolsFileURL(c *gc.C) {
	storageName := tools.StorageName(version.Current)
	sftpURL, err := s.env.Storage().URL(storageName)
	c.Assert(err, gc.IsNil)
	fileURL := fmt.Sprintf("file://%s/%s", s.env.storageDir, storageName)
	s.testBootstrapToolsURL(c, sftpURL, fileURL)
}

func (s *bootstrapSuite) TestBootstrapToolsExternalURL(c *gc.C) {
	const externalURL = "http://test.invalid/tools.tgz"
	s.testBootstrapToolsURL(c, externalURL, externalURL)
}

func (s *bootstrapSuite) testBootstrapToolsURL(c *gc.C, toolsURL, expectedURL string) {
	s.PatchValue(manual.ProvisionMachineAgent, func(host string, mcfg *cloudinit.MachineConfig, w io.Writer) error {
		c.Assert(mcfg.Tools.URL, gc.Equals, expectedURL)
		return nil
	})
	args := s.getArgs(c)
	args.PossibleTools[0].URL = toolsURL
	defer fakeSSH{SkipDetection: true}.install(c).Restore()
	err := manual.Bootstrap(args)
	c.Assert(err, gc.IsNil)
}
<|endoftext|>