<commit_before>// Copyright 2013 The go-github AUTHORS. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package github

import (
	"context"
	"fmt"
)

// Membership represents the status of a user's membership in an organization or team.
type Membership struct {
	URL *string `json:"url,omitempty"`

	// State is the user's status within the organization or team.
	// Possible values are: "active", "pending"
	State *string `json:"state,omitempty"`

	// Role identifies the user's role within the organization or team.
	// Possible values for organization membership:
	//     member - non-owner organization member
	//     admin - organization owner
	//
	// Possible values for team membership are:
	//     member - a normal member of the team
	//     maintainer - a team maintainer. Able to add/remove other team
	//                  members, promote other team members to team
	//                  maintainer, and edit the team’s name and description
	Role *string `json:"role,omitempty"`

	// For organization membership, the API URL of the organization.
	OrganizationURL *string `json:"organization_url,omitempty"`

	// For organization membership, the organization the membership is for.
	Organization *Organization `json:"organization,omitempty"`

	// For organization membership, the user the membership is for.
	User *User `json:"user,omitempty"`
}

func (m Membership) String() string {
	return Stringify(m)
}

// ListMembersOptions specifies optional parameters to the
// OrganizationsService.ListMembers method.
type ListMembersOptions struct {
	// If true (or if the authenticated user is not an owner of the
	// organization), list only publicly visible members.
	PublicOnly bool `url:"-"`

	// Filter members returned in the list. Possible values are:
	// 2fa_disabled, all. Default is "all".
	Filter string `url:"filter,omitempty"`

	// Role filters members returned by their role in the organization.
	// Possible values are:
	//     all - all members of the organization, regardless of role
	//     admin - organization owners
	//     member - non-organization members
	//
	// Default is "all".
	Role string `url:"role,omitempty"`

	ListOptions
}

// ListMembers lists the members for an organization. If the authenticated
// user is an owner of the organization, this will return both concealed and
// public members, otherwise it will only return public members.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#members-list
func (s *OrganizationsService) ListMembers(ctx context.Context, org string, opt *ListMembersOptions) ([]*User, *Response, error) {
	var u string
	if opt != nil && opt.PublicOnly {
		u = fmt.Sprintf("orgs/%v/public_members", org)
	} else {
		u = fmt.Sprintf("orgs/%v/members", org)
	}
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	var members []*User
	resp, err := s.client.Do(ctx, req, &members)
	if err != nil {
		return nil, resp, err
	}

	return members, resp, nil
}

// IsMember checks if a user is a member of an organization.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#check-membership
func (s *OrganizationsService) IsMember(ctx context.Context, org, user string) (bool, *Response, error) {
	u := fmt.Sprintf("orgs/%v/members/%v", org, user)
	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return false, nil, err
	}

	resp, err := s.client.Do(ctx, req, nil)
	member, err := parseBoolResponse(err)
	return member, resp, err
}

// IsPublicMember checks if a user is a public member of an organization.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#check-public-membership
func (s *OrganizationsService) IsPublicMember(ctx context.Context, org, user string) (bool, *Response, error) {
	u := fmt.Sprintf("orgs/%v/public_members/%v", org, user)
	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return false, nil, err
	}

	resp, err := s.client.Do(ctx, req, nil)
	member, err := parseBoolResponse(err)
	return member, resp, err
}

// RemoveMember removes a user from all teams of an organization.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#remove-a-member
func (s *OrganizationsService) RemoveMember(ctx context.Context, org, user string) (*Response, error) {
	u := fmt.Sprintf("orgs/%v/members/%v", org, user)
	req, err := s.client.NewRequest("DELETE", u, nil)
	if err != nil {
		return nil, err
	}

	return s.client.Do(ctx, req, nil)
}

// PublicizeMembership publicizes a user's membership in an organization. (A
// user cannot publicize the membership for another user.)
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#publicize-a-users-membership
func (s *OrganizationsService) PublicizeMembership(ctx context.Context, org, user string) (*Response, error) {
	u := fmt.Sprintf("orgs/%v/public_members/%v", org, user)
	req, err := s.client.NewRequest("PUT", u, nil)
	if err != nil {
		return nil, err
	}

	return s.client.Do(ctx, req, nil)
}

// ConcealMembership conceals a user's membership in an organization.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#conceal-a-users-membership
func (s *OrganizationsService) ConcealMembership(ctx context.Context, org, user string) (*Response, error) {
	u := fmt.Sprintf("orgs/%v/public_members/%v", org, user)
	req, err := s.client.NewRequest("DELETE", u, nil)
	if err != nil {
		return nil, err
	}

	return s.client.Do(ctx, req, nil)
}

// ListOrgMembershipsOptions specifies optional parameters to the
// OrganizationsService.ListOrgMemberships method.
type ListOrgMembershipsOptions struct {
	// Filter memberships to include only those with the specified state.
	// Possible values are: "active", "pending".
	State string `url:"state,omitempty"`

	ListOptions
}

// ListOrgMemberships lists the organization memberships for the authenticated user.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#list-your-organization-memberships
func (s *OrganizationsService) ListOrgMemberships(ctx context.Context, opt *ListOrgMembershipsOptions) ([]*Membership, *Response, error) {
	u := "user/memberships/orgs"
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	var memberships []*Membership
	resp, err := s.client.Do(ctx, req, &memberships)
	if err != nil {
		return nil, resp, err
	}

	return memberships, resp, nil
}

// GetOrgMembership gets the membership for a user in a specified organization.
// Passing an empty string for user will get the membership for the
// authenticated user.
//
// GitHub API docs:
// https://developer.github.com/v3/orgs/members/#get-organization-membership
// https://developer.github.com/v3/orgs/members/#get-your-organization-membership
func (s *OrganizationsService) GetOrgMembership(ctx context.Context, user, org string) (*Membership, *Response, error) {
	var u string
	if user != "" {
		u = fmt.Sprintf("orgs/%v/memberships/%v", org, user)
	} else {
		u = fmt.Sprintf("user/memberships/orgs/%v", org)
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	membership := new(Membership)
	resp, err := s.client.Do(ctx, req, membership)
	if err != nil {
		return nil, resp, err
	}

	return membership, resp, nil
}

// EditOrgMembership edits the membership for user in specified organization.
// Passing an empty string for user will edit the membership for the
// authenticated user.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#add-or-update-organization-membership
// GitHub API docs: https://developer.github.com/v3/orgs/members/#edit-your-organization-membership
func (s *OrganizationsService) EditOrgMembership(ctx context.Context, user, org string, membership *Membership) (*Membership, *Response, error) {
	var u, method string
	if user != "" {
		u = fmt.Sprintf("orgs/%v/memberships/%v", org, user)
		method = "PUT"
	} else {
		u = fmt.Sprintf("user/memberships/orgs/%v", org)
		method = "PATCH"
	}

	req, err := s.client.NewRequest(method, u, membership)
	if err != nil {
		return nil, nil, err
	}

	m := new(Membership)
	resp, err := s.client.Do(ctx, req, m)
	if err != nil {
		return nil, resp, err
	}

	return m, resp, nil
}

// RemoveOrgMembership removes user from the specified organization. If the
// user has been invited to the organization, this will cancel their invitation.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#remove-organization-membership
func (s *OrganizationsService) RemoveOrgMembership(ctx context.Context, user, org string) (*Response, error) {
	u := fmt.Sprintf("orgs/%v/memberships/%v", org, user)
	req, err := s.client.NewRequest("DELETE", u, nil)
	if err != nil {
		return nil, err
	}

	return s.client.Do(ctx, req, nil)
}

// ListPendingOrgInvitations returns a list of pending invitations.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#list-pending-organization-invitations
func (s *OrganizationsService) ListPendingOrgInvitations(ctx context.Context, org string, opt *ListOptions) ([]*Invitation, *Response, error) {
	u := fmt.Sprintf("orgs/%v/invitations", org)
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	var pendingInvitations []*Invitation
	resp, err := s.client.Do(ctx, req, &pendingInvitations)
	if err != nil {
		return nil, resp, err
	}
	return pendingInvitations, resp, nil
}

// CreateOrgInvitationOptions specifies the parameters to the OrganizationService.Invite
// method.
type CreateOrgInvitationOptions struct {
	// GitHub user ID for the person you are inviting. Not required if you provide Email.
	InviteeID *int64 `json:"invitee_id,omitempty"`
	// Email address of the person you are inviting, which can be an existing GitHub user.
	// Not required if you provide InviteeID
	Email *string `json:"email,omitempty"`
	// Specify role for new member. Can be one of:
	// * admin - Organization owners with full administrative rights to the
	// 	 organization and complete access to all repositories and teams.
	// * direct_member - Non-owner organization members with ability to see
	//   other members and join teams by invitation.
	// * billing_manager - Non-owner organization members with ability to
	//   manage the billing settings of your organization.
	// Default is "direct_member".
	Role   *string `json:"role"`
	TeamID []int64 `json:"team_ids"`
}

// CreateOrgInvitation invites people to an organization by using their GitHub user ID or their email address.
// In order to create invitations in an organization,
// the authenticated user must be an organization owner.
//
// https://developer.github.com/v3/orgs/members/#create-organization-invitation
func (s *OrganizationsService) CreateOrgInvitation(ctx context.Context, org string, opt *CreateOrgInvitationOptions) (*Invitation, *Response, error) {
	u := fmt.Sprintf("orgs/%v/invitations", org)

	req, err := s.client.NewRequest("POST", u, opt)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches.
	req.Header.Set("Accept", mediaTypeOrganizationInvitationPreview)

	var invitation *Invitation
	resp, err := s.client.Do(ctx, req, &invitation)
	if err != nil {
		return nil, resp, err
	}
	return invitation, resp, nil
}

// ListOrgInvitationTeams lists all teams associated with an invitation. In order to see invitations in an organization,
// the authenticated user must be an organization owner.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#list-organization-invitation-teams
func (s *OrganizationsService) ListOrgInvitationTeams(ctx context.Context, org, invitationID string, opt *ListOptions) ([]*Team, *Response, error) {
	u := fmt.Sprintf("orgs/%v/invitations/%v/teams", org, invitationID)
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches.
	req.Header.Set("Accept", mediaTypeOrganizationInvitationPreview)

	var orgInvitationTeams []*Team
	resp, err := s.client.Do(ctx, req, &orgInvitationTeams)
	if err != nil {
		return nil, resp, err
	}
	return orgInvitationTeams, resp, nil
}
<commit_msg>Update orgs_members.go (#871)<commit_after>// Copyright 2013 The go-github AUTHORS. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package github

import (
	"context"
	"fmt"
)

// Membership represents the status of a user's membership in an organization or team.
type Membership struct {
	URL *string `json:"url,omitempty"`

	// State is the user's status within the organization or team.
	// Possible values are: "active", "pending"
	State *string `json:"state,omitempty"`

	// Role identifies the user's role within the organization or team.
	// Possible values for organization membership:
	//     member - non-owner organization member
	//     admin - organization owner
	//
	// Possible values for team membership are:
	//     member - a normal member of the team
	//     maintainer - a team maintainer. Able to add/remove other team
	//                  members, promote other team members to team
	//                  maintainer, and edit the team’s name and description
	Role *string `json:"role,omitempty"`

	// For organization membership, the API URL of the organization.
	OrganizationURL *string `json:"organization_url,omitempty"`

	// For organization membership, the organization the membership is for.
	Organization *Organization `json:"organization,omitempty"`

	// For organization membership, the user the membership is for.
	User *User `json:"user,omitempty"`
}

func (m Membership) String() string {
	return Stringify(m)
}

// ListMembersOptions specifies optional parameters to the
// OrganizationsService.ListMembers method.
type ListMembersOptions struct {
	// If true (or if the authenticated user is not an owner of the
	// organization), list only publicly visible members.
	PublicOnly bool `url:"-"`

	// Filter members returned in the list. Possible values are:
	// 2fa_disabled, all. Default is "all".
	Filter string `url:"filter,omitempty"`

	// Role filters members returned by their role in the organization.
	// Possible values are:
	//     all - all members of the organization, regardless of role
	//     admin - organization owners
	//     member - non-owner organization members
	//
	// Default is "all".
	Role string `url:"role,omitempty"`

	ListOptions
}

// ListMembers lists the members for an organization. If the authenticated
// user is an owner of the organization, this will return both concealed and
// public members, otherwise it will only return public members.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#members-list
func (s *OrganizationsService) ListMembers(ctx context.Context, org string, opt *ListMembersOptions) ([]*User, *Response, error) {
	var u string
	if opt != nil && opt.PublicOnly {
		u = fmt.Sprintf("orgs/%v/public_members", org)
	} else {
		u = fmt.Sprintf("orgs/%v/members", org)
	}
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	var members []*User
	resp, err := s.client.Do(ctx, req, &members)
	if err != nil {
		return nil, resp, err
	}

	return members, resp, nil
}

// IsMember checks if a user is a member of an organization.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#check-membership
func (s *OrganizationsService) IsMember(ctx context.Context, org, user string) (bool, *Response, error) {
	u := fmt.Sprintf("orgs/%v/members/%v", org, user)
	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return false, nil, err
	}

	resp, err := s.client.Do(ctx, req, nil)
	member, err := parseBoolResponse(err)
	return member, resp, err
}

// IsPublicMember checks if a user is a public member of an organization.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#check-public-membership
func (s *OrganizationsService) IsPublicMember(ctx context.Context, org, user string) (bool, *Response, error) {
	u := fmt.Sprintf("orgs/%v/public_members/%v", org, user)
	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return false, nil, err
	}

	resp, err := s.client.Do(ctx, req, nil)
	member, err := parseBoolResponse(err)
	return member, resp, err
}

// RemoveMember removes a user from all teams of an organization.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#remove-a-member
func (s *OrganizationsService) RemoveMember(ctx context.Context, org, user string) (*Response, error) {
	u := fmt.Sprintf("orgs/%v/members/%v", org, user)
	req, err := s.client.NewRequest("DELETE", u, nil)
	if err != nil {
		return nil, err
	}

	return s.client.Do(ctx, req, nil)
}

// PublicizeMembership publicizes a user's membership in an organization. (A
// user cannot publicize the membership for another user.)
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#publicize-a-users-membership
func (s *OrganizationsService) PublicizeMembership(ctx context.Context, org, user string) (*Response, error) {
	u := fmt.Sprintf("orgs/%v/public_members/%v", org, user)
	req, err := s.client.NewRequest("PUT", u, nil)
	if err != nil {
		return nil, err
	}

	return s.client.Do(ctx, req, nil)
}

// ConcealMembership conceals a user's membership in an organization.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#conceal-a-users-membership
func (s *OrganizationsService) ConcealMembership(ctx context.Context, org, user string) (*Response, error) {
	u := fmt.Sprintf("orgs/%v/public_members/%v", org, user)
	req, err := s.client.NewRequest("DELETE", u, nil)
	if err != nil {
		return nil, err
	}

	return s.client.Do(ctx, req, nil)
}

// ListOrgMembershipsOptions specifies optional parameters to the
// OrganizationsService.ListOrgMemberships method.
type ListOrgMembershipsOptions struct {
	// Filter memberships to include only those with the specified state.
	// Possible values are: "active", "pending".
	State string `url:"state,omitempty"`

	ListOptions
}

// ListOrgMemberships lists the organization memberships for the authenticated user.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#list-your-organization-memberships
func (s *OrganizationsService) ListOrgMemberships(ctx context.Context, opt *ListOrgMembershipsOptions) ([]*Membership, *Response, error) {
	u := "user/memberships/orgs"
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	var memberships []*Membership
	resp, err := s.client.Do(ctx, req, &memberships)
	if err != nil {
		return nil, resp, err
	}

	return memberships, resp, nil
}

// GetOrgMembership gets the membership for a user in a specified organization.
// Passing an empty string for user will get the membership for the
// authenticated user.
//
// GitHub API docs:
// https://developer.github.com/v3/orgs/members/#get-organization-membership
// https://developer.github.com/v3/orgs/members/#get-your-organization-membership
func (s *OrganizationsService) GetOrgMembership(ctx context.Context, user, org string) (*Membership, *Response, error) {
	var u string
	if user != "" {
		u = fmt.Sprintf("orgs/%v/memberships/%v", org, user)
	} else {
		u = fmt.Sprintf("user/memberships/orgs/%v", org)
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	membership := new(Membership)
	resp, err := s.client.Do(ctx, req, membership)
	if err != nil {
		return nil, resp, err
	}

	return membership, resp, nil
}

// EditOrgMembership edits the membership for user in specified organization.
// Passing an empty string for user will edit the membership for the
// authenticated user.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#add-or-update-organization-membership
// GitHub API docs: https://developer.github.com/v3/orgs/members/#edit-your-organization-membership
func (s *OrganizationsService) EditOrgMembership(ctx context.Context, user, org string, membership *Membership) (*Membership, *Response, error) {
	var u, method string
	if user != "" {
		u = fmt.Sprintf("orgs/%v/memberships/%v", org, user)
		method = "PUT"
	} else {
		u = fmt.Sprintf("user/memberships/orgs/%v", org)
		method = "PATCH"
	}

	req, err := s.client.NewRequest(method, u, membership)
	if err != nil {
		return nil, nil, err
	}

	m := new(Membership)
	resp, err := s.client.Do(ctx, req, m)
	if err != nil {
		return nil, resp, err
	}

	return m, resp, nil
}

// RemoveOrgMembership removes user from the specified organization. If the
// user has been invited to the organization, this will cancel their invitation.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#remove-organization-membership
func (s *OrganizationsService) RemoveOrgMembership(ctx context.Context, user, org string) (*Response, error) {
	u := fmt.Sprintf("orgs/%v/memberships/%v", org, user)
	req, err := s.client.NewRequest("DELETE", u, nil)
	if err != nil {
		return nil, err
	}

	return s.client.Do(ctx, req, nil)
}

// ListPendingOrgInvitations returns a list of pending invitations.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#list-pending-organization-invitations
func (s *OrganizationsService) ListPendingOrgInvitations(ctx context.Context, org string, opt *ListOptions) ([]*Invitation, *Response, error) {
	u := fmt.Sprintf("orgs/%v/invitations", org)
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	var pendingInvitations []*Invitation
	resp, err := s.client.Do(ctx, req, &pendingInvitations)
	if err != nil {
		return nil, resp, err
	}
	return pendingInvitations, resp, nil
}

// CreateOrgInvitationOptions specifies the parameters to the OrganizationService.Invite
// method.
type CreateOrgInvitationOptions struct {
	// GitHub user ID for the person you are inviting. Not required if you provide Email.
	InviteeID *int64 `json:"invitee_id,omitempty"`
	// Email address of the person you are inviting, which can be an existing GitHub user.
	// Not required if you provide InviteeID
	Email *string `json:"email,omitempty"`
	// Specify role for new member. Can be one of:
	// * admin - Organization owners with full administrative rights to the
	// 	 organization and complete access to all repositories and teams.
	// * direct_member - Non-owner organization members with ability to see
	//   other members and join teams by invitation.
	// * billing_manager - Non-owner organization members with ability to
	//   manage the billing settings of your organization.
	// Default is "direct_member".
	Role   *string `json:"role"`
	TeamID []int64 `json:"team_ids"`
}

// CreateOrgInvitation invites people to an organization by using their GitHub user ID or their email address.
// In order to create invitations in an organization,
// the authenticated user must be an organization owner.
//
// https://developer.github.com/v3/orgs/members/#create-organization-invitation
func (s *OrganizationsService) CreateOrgInvitation(ctx context.Context, org string, opt *CreateOrgInvitationOptions) (*Invitation, *Response, error) {
	u := fmt.Sprintf("orgs/%v/invitations", org)

	req, err := s.client.NewRequest("POST", u, opt)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches.
	req.Header.Set("Accept", mediaTypeOrganizationInvitationPreview)

	var invitation *Invitation
	resp, err := s.client.Do(ctx, req, &invitation)
	if err != nil {
		return nil, resp, err
	}
	return invitation, resp, nil
}

// ListOrgInvitationTeams lists all teams associated with an invitation. In order to see invitations in an organization,
// the authenticated user must be an organization owner.
//
// GitHub API docs: https://developer.github.com/v3/orgs/members/#list-organization-invitation-teams
func (s *OrganizationsService) ListOrgInvitationTeams(ctx context.Context, org, invitationID string, opt *ListOptions) ([]*Team, *Response, error) {
	u := fmt.Sprintf("orgs/%v/invitations/%v/teams", org, invitationID)
	u, err := addOptions(u, opt)
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil)
	if err != nil {
		return nil, nil, err
	}

	// TODO: remove custom Accept header when this API fully launches.
	req.Header.Set("Accept", mediaTypeOrganizationInvitationPreview)

	var orgInvitationTeams []*Team
	resp, err := s.client.Do(ctx, req, &orgInvitationTeams)
	if err != nil {
		return nil, resp, err
	}
	return orgInvitationTeams, resp, nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"os"
	"syscall"
)

type requestForVotesResponse struct {
	term           int
	hasGrantedVote bool
}

type voteResponse struct {
	serverIndex int
	resp        requestForVotesResponse
}

type voteRequest struct {
	ip     string
	termID int
}

type term struct {
	id       int
	votes    int
	votedFor string
}

func (s *server) sendRequestForVotes(receiver string, respChan chan RequestForVotesResponse) {
	v := url.Values{}
	v.set("candidateID", s.id)
	v.set("term", s.term)
	v.set("lastLogIndex", len(s.db.log))
	v.set("lastLogTerm", s.db.log[len(s.db.log)-1].term)
	resp, err := http.PostForm(server+"/votes", v)
	if err != nil {
		fmt.Println("Couldn't send request for votes to " + server)
	}
	defer resp.Body.Close()
	r = &RequestForVotesResponse{}
	json.NewDecoder(resp.Body).Decode(r)
	v = &VoteResponse{receiver, r}
	respChan <- v
}

func (server *server) startElection() {
	server.state = "candidate"
	server.term += 1
	server.term.votes += 1
	respChan = make(chan voteResponse)
	for receiverIndex, _ := range server.config {
		go server.sendRequestForVotes(receiverIndex)
	}
	voteCount = 0
	for {
		vote := <-respChan
		if voteResponse.resp.hasGrantedVote {
			voteCount++
		}
		if voteCount > len(s.config)/2 {
			s.state = "leader"
			break
		}
	}

}

func (s *server) handleRequestForVote(request voteRequest, w http.ResponseWriter) {
	if request.termID < s.term.id {
		fmt.Fprint(w, false)
	} else {
		cond1 := s.term.vote == ""
		cond2 := s.term.vote == request.candidateID
		cond3 := request.lastLogIndex >= s.lastLogIndex
		if (cond1 || cond2) && cond3 {
			s.electionTimeout.resetTimeout()
			s.term.vote = request.candidateID
			s.term.id = request.termID
			fmt.Fprint(w, true)
		}
	}
}
<commit_msg>end election if all servers have replied, remove unused types<commit_after>package main

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"os"
	"syscall"
)

type requestForVotesResponse struct {
	term           int
	hasGrantedVote bool
}

type voteResponse struct {
	serverIndex int
	resp        requestForVotesResponse
}

func (s *server) handleRequestForVote(request voteRequest, w http.ResponseWriter) {
	if request.term < s.term {
		fmt.Fprint(w, false)
	} else {
		cond1 := s.votedFor == ""
		cond2 := s.votedFor == request.candidateID
		cond3 := request.lastLogIndex >= len(s.log)
		if (cond1 || cond2) && cond3 {
			s.electionTimeout.resetTimeout()
			s.votedFor = request.candidateID
			s.term = request.term
			fmt.Fprint(w, true)
		}
	}
}

func (s *server) sendRequestForVote(receiver string, respChan chan RequestForVoteResponse) {
	v := url.Values{}
	v.set("candidateID", s.id)
	v.set("term", s.term)
	v.set("lastLogIndex", len(s.db.log))
	v.set("lastLogTerm", s.db.log[len(s.db.log)-1].term)
	resp, err := http.PostForm(server+"/votes", v)
	if err != nil {
		fmt.Println("Couldn't send request for votes to " + server)
	}
	defer resp.Body.Close()
	r = &RequestForVoteResponse{}
	json.NewDecoder(resp.Body).Decode(r)
	v = &VoteResponse{receiver, r}
	respChan <- v
}

func (server *server) startElection() {
	server.state = "candidate"
	server.term += 1
	server.term.votes += 1
	respChan = make(chan voteResponse)
	for receiverIndex, _ := range server.config {
		go server.sendRequestForVotes(receiverIndex)
	}
	voteCount = 0
	responseCount = 0
	for {
		vote := <-respChan
		responseCount++
		if vote.resp.hasGrantedVote {
			voteCount++
		}
		if voteCount > len(s.config)/2 {
			s.state = "leader"
			break
		}
		if responseCount == len(s.config) {
			break
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net/http"

	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/trusch/jwtd/jwt"
)

type Proxy struct {
	cfg     *Config
	proxies map[string]*SingleProxy
	router  *mux.Router
	key     interface{}
}

func NewProxy(cfg *Config) (*Proxy, error) {
	r := mux.NewRouter()
	key, err := jwt.LoadPublicKey(cfg.Cert)
	if err != nil {
		return nil, err
	}
	proxy := &Proxy{cfg: cfg, key: key}
	for host, hostCfg := range cfg.Hosts {
		singleProxy, err := NewSingleProxy(hostCfg.Project, host, hostCfg.Backend, hostCfg.Routes, proxy.key)
		if err != nil {
			return nil, err
		}
		var handler http.Handler = singleProxy
		if hostCfg.CORS != nil {
			headers := handlers.AllowedHeaders(hostCfg.CORS.AllowedHeaders)
			origins := handlers.AllowedOrigins(hostCfg.CORS.AllowedOrigins)
			methods := handlers.AllowedMethods(hostCfg.CORS.AllowedMethods)
			corsWrapper := handlers.CORS(headers, origins, methods)
			handler = corsWrapper(handler)
		}
		r.Host(host).Handler(handler)
	}
	proxy.router = r
	return proxy, nil
}

func (proxy *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Print("incoming request for host: ", r.Host)
	proxy.router.ServeHTTP(w, r)
}
<commit_msg>added logging to CORS handling;<commit_after>package main

import (
	"log"
	"net/http"

	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/trusch/jwtd/jwt"
)

type Proxy struct {
	cfg     *Config
	proxies map[string]*SingleProxy
	router  *mux.Router
	key     interface{}
}

func NewProxy(cfg *Config) (*Proxy, error) {
	r := mux.NewRouter()
	key, err := jwt.LoadPublicKey(cfg.Cert)
	if err != nil {
		return nil, err
	}
	proxy := &Proxy{cfg: cfg, key: key}
	for host, hostCfg := range cfg.Hosts {
		singleProxy, err := NewSingleProxy(hostCfg.Project, host, hostCfg.Backend, hostCfg.Routes, proxy.key)
		if err != nil {
			return nil, err
		}
		var handler http.Handler = singleProxy
		if hostCfg.CORS != nil {
			log.Printf("use CORS for %v", host)
			log.Printf("allowed Headers: %v", hostCfg.CORS.AllowedHeaders)
			log.Printf("allowed Origins: %v", hostCfg.CORS.AllowedOrigins)
			log.Printf("allowed Methods: %v", hostCfg.CORS.AllowedMethods)
			headers := handlers.AllowedHeaders(hostCfg.CORS.AllowedHeaders)
			origins := handlers.AllowedOrigins(hostCfg.CORS.AllowedOrigins)
			methods := handlers.AllowedMethods(hostCfg.CORS.AllowedMethods)
			corsWrapper := handlers.CORS(headers, origins, methods)
			handler = corsWrapper(handler)
		}
		r.Host(host).Handler(handler)
	}
	proxy.router = r
	return proxy, nil
}

func (proxy *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.Print("incoming request for host: ", r.Host)
	proxy.router.ServeHTTP(w, r)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"sync"

	"time"

	"encoding/json"

	"github.com/jmoiron/sqlx"
	"github.com/vincentserpoul/playwithsql/bench"
	"github.com/vincentserpoul/playwithsql/dbhandler"
	"github.com/vincentserpoul/playwithsql/status"
)

// Number of retries after query returns an error
const maxRetryCount = 3

// Results to be returned
type Results struct {
	DBType       string
	MaxConns     int
	Date         time.Time
	BenchResults []BenchResult
}

// BenchResult data
type BenchResult struct {
	Action      string
	Loops       int
	Concurrency int
	PauseTime   time.Duration
	Errors      int
	Min         time.Duration
	Max         time.Duration
	Median      time.Duration
	StandDev    time.Duration
	Throughput  int
}

func main() {

	// Flags
	dbName := "playwithsql"
	dbType := flag.String("db", "mysql", "type of db to bench: mysql, cockroachdb, postgres...")
	schemaType := flag.String("sch", "islatest", "type of schema to use, is latest, latest status...")
	dbHost := flag.String("host", "127.0.0.1", "host IP")
	loops := flag.Int("loops", 100, "number of loops")
	maxConns := flag.Int("maxconns", 10, "number of max connections")
	concurrency := flag.Int("concurrency", 100, "number of go routines created")

	flag.Parse()

	if *concurrency > *loops {
		*concurrency = *loops
	}

	db, err := dbhandler.Get(*dbType, *dbHost, dbName)
	if err != nil {
		log.Fatalf("%s - %s - %s, \n%v", *dbType, *dbHost, dbName, err)
	}

	// Connection
	islatestSQLLink := status.GetSQLIntImpl(*dbType, *schemaType)

	ctx := context.Background()
	err = islatestSQLLink.MigrateDown(ctx, db)
	if err != nil {
		log.Fatalf("%v", err)
	}
	err = islatestSQLLink.MigrateUp(ctx, db)
	if err != nil {
		log.Fatalf("%v", err)
	}

	// Number of max connections
	// TODO set the param in the db config
	db.SetMaxOpenConns(*maxConns)
	db.SetMaxIdleConns(*maxConns)

	var results = Results{
		DBType:   *dbType,
		MaxConns: *maxConns,
		Date:     time.Now(),
	}

	// Create
	createResults, testEntityoneIDs, err := BenchmarkCreate(ctx, *loops, *concurrency, db, islatestSQLLink)
	if err != nil {
		log.Fatalf("%v", err)
	}
	results.BenchResults = append(results.BenchResults, createResults)

	// Update
	updateResults, err := BenchmarkUpdateStatus(ctx, *loops, *concurrency, db, islatestSQLLink, testEntityoneIDs)
	if err != nil {
		log.Fatalf("%v", err)
	}
	results.BenchResults = append(results.BenchResults, updateResults)

	// Select by status
	selectByStatusResults, err := BenchmarkSelectEntityoneByStatus(ctx, *loops, *concurrency, db, islatestSQLLink)
	if err != nil {
		log.Fatalf("%v", err)
	}
	results.BenchResults = append(results.BenchResults, selectByStatusResults)

	// Select by PK
	selectByPKResults, err := BenchmarkSelectEntityoneOneByPK(ctx, *loops, *concurrency, db, islatestSQLLink, testEntityoneIDs)
	if err != nil {
		log.Fatalf("%v", err)
	}
	results.BenchResults = append(results.BenchResults, selectByPKResults)

	jsonResults, err := json.Marshal(results)
	if err != nil {
		log.Fatalf("%v", err)
	}

	fmt.Printf("%s\n", jsonResults)
}

// BenchmarkCreate will loop a loops number of time and give the resulting time taken
func BenchmarkCreate(
	ctx context.Context,
	loops int,
	concurrency int,
	dbConn *sqlx.DB,
	benchSQLLink *status.SQLIntImpl,
) (
	results BenchResult,
	testEntityoneIDs []int64,
	err error,
) {
	entityIDsC := make(chan int64)
	latenciesC, errorC, doneC, resultsC := handleResults()

	before := time.Now()
	var wg sync.WaitGroup

	// Pause time
	dynPauseTime := 1 * time.Millisecond
	dynPauseTimeC := dynPauseTimeInit(&dynPauseTime)
	defer close(dynPauseTimeC)

	// Launch as many concurrent connections as asked
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func(routineNum int, ctx context.Context, wg *sync.WaitGroup) {
			for j := 0; j < loops/concurrency; j++ {
				time.Sleep(dynPauseTime)
				var e status.Entityone
				beforeLocal := time.Now()
				ok := false
				var errCr error
				retryCount := 0
				for retryCount < maxRetryCount && !ok {
					// Timeout
					sqlCtx, sqlCncl := context.WithTimeout(ctx, 250*time.Millisecond)
					defer sqlCncl()

					// For each error, we add some pause time
					errCr = e.Create(sqlCtx, dbConn, benchSQLLink)
					if errCr != nil {
						retryCount++
						time.Sleep(dynPauseTime)
						dynPauseTimeC <- 1 * time.Millisecond
					} else {
						ok = true
					}
				}
				if errCr != nil {
					errorC <- errCr
				} else {
					latenciesC <- time.Since(beforeLocal)
					entityIDsC <- e.ID
					// If no error, we increment down a little bit
					dynPauseTimeC <- -1 * time.Millisecond
				}
			}
			wg.Done()
		}(i, ctx, &wg)
	}

	// Receive the entityIDs
	go func() {
		for entityID := range entityIDsC {
			testEntityoneIDs = append(testEntityoneIDs, entityID)
		}
	}()

	wg.Wait()
	doneC <- true
	rawRes := <-resultsC
	timeTaken := time.Since(before)
	return BenchResult{
			Action:      "create",
			Loops:       loops,
			Concurrency: concurrency,
			PauseTime:   dynPauseTime,
			Errors:      rawRes.errCount,
			Min:         bench.GetMin(rawRes.latencies),
			Max:         bench.GetMax(rawRes.latencies),
			Median:      bench.GetMedian(rawRes.latencies),
			StandDev:    bench.GetStandardDeviation(rawRes.latencies),
			Throughput:  int(float64(loops) / timeTaken.Seconds()),
		},
		testEntityoneIDs,
		nil
}

// BenchmarkUpdateStatus benchmark for status updates (include deletes)
func BenchmarkUpdateStatus(
	ctx context.Context,
	loops int,
	concurrency int,
	dbConn *sqlx.DB,
	benchSQLLink *status.SQLIntImpl,
	testEntityoneIDs []int64,
) (
	results BenchResult,
	err error,
) {
	if len(testEntityoneIDs) == 0 {
		return results, fmt.Errorf("BenchmarkUpdateStatus: no entity created, nothing to update")
	}

	latenciesC, errorC, doneC, resultsC := handleResults()

	before := time.Now()
	var wg sync.WaitGroup

	// Pause time
	dynPauseTime := 1 * time.Millisecond
	dynPauseTimeC := dynPauseTimeInit(&dynPauseTime)
	defer close(dynPauseTimeC)

	for i := 0; i < concurrency; i++ {
		wg.Add(1)

		go func(ctx context.Context, wg *sync.WaitGroup) {
			for j := 0; j < loops/concurrency; j++ {
				time.Sleep(dynPauseTime)
				var e status.Entityone
				e.ID = testEntityoneIDs[i%len(testEntityoneIDs)]
				beforeLocal := time.Now()
				ok := false
				var errU error
				retryCount := 0
				for retryCount < maxRetryCount && !ok {
					// Timeout
					sqlCtx, sqlCncl := context.WithTimeout(ctx, 250*time.Millisecond)
					defer sqlCncl()
					errU = e.UpdateStatus(sqlCtx, dbConn, benchSQLLink, status.ActionCancel, status.StatusCancelled)
					if errU != nil {
						retryCount++
						time.Sleep(dynPauseTime)
						dynPauseTimeC <- 1 * time.Millisecond
					} else {
						ok = true
					}
				}
				if errU != nil {
					errorC <- errU
				} else {
					latenciesC <- time.Since(beforeLocal)
					// If no error, we increment down a little bit
					dynPauseTimeC <- -1 * time.Millisecond
				}
			}
			wg.Done()
		}(ctx, &wg)
	}

	wg.Wait()
	doneC <- true
	rawRes := <-resultsC
	timeTaken := time.Since(before)

	return BenchResult{
			Action:      "updateStatus",
			Loops:       loops,
			Concurrency: concurrency,
			PauseTime:   dynPauseTime,
			Errors:      rawRes.errCount,
			Min:         bench.GetMin(rawRes.latencies),
			Max:         bench.GetMax(rawRes.latencies),
			Median:      bench.GetMedian(rawRes.latencies),
			StandDev:    bench.GetStandardDeviation(rawRes.latencies),
			Throughput:  int(float64(loops) / timeTaken.Seconds()),
		},
		nil

}

// BenchmarkSelectEntityoneByStatus benchmark with select by status
func BenchmarkSelectEntityoneByStatus(
	ctx context.Context,
	loops int,
	concurrency int,
	dbConn *sqlx.DB,
	benchSQLLink *status.SQLIntImpl,
) (
	results BenchResult,
	err error,
) {
	latenciesC, errorC, doneC, resultsC := handleResults()

	var wg sync.WaitGroup
	before := time.Now()

	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func(ctx context.Context, wg *sync.WaitGroup) {
			for j := 0; j < loops/concurrency; j++ {
				beforeLocal := time.Now()
				sqlCtx, sqlCncl := context.WithTimeout(ctx, 100*time.Millisecond)
				defer sqlCncl()
				_, errSel := status.SelectEntityoneByStatus(sqlCtx, dbConn, benchSQLLink, status.StatusCancelled)
				if errSel != nil {
					errorC <- errSel
				} else {
					latenciesC <- time.Since(beforeLocal)
				}
			}
			wg.Done()
		}(ctx, &wg)
	}

	wg.Wait()
	doneC <- true
	rawRes := <-resultsC
	timeTaken := time.Since(before)

	return BenchResult{
			Action:      "selectEntityoneByStatus",
			Loops:       loops,
			Concurrency: concurrency,
			PauseTime:   0,
			Errors:      rawRes.errCount,
			Min:         bench.GetMin(rawRes.latencies),
			Max:         bench.GetMax(rawRes.latencies),
			Median:      bench.GetMedian(rawRes.latencies),
			StandDev:    bench.GetStandardDeviation(rawRes.latencies),
			Throughput:  int(float64(loops) / timeTaken.Seconds()),
		},
		nil
}

// BenchmarkSelectEntityoneOneByPK benchmark with select by primary key
func BenchmarkSelectEntityoneOneByPK(
	ctx context.Context,
	loops int,
	concurrency int,
	dbConn *sqlx.DB,
	benchSQLLink *status.SQLIntImpl,
	testEntityoneIDs []int64,
) (
	results BenchResult,
	err error,
) {
	latenciesC, errorC, doneC, resultsC := handleResults()

	before := time.Now()
	var wg sync.WaitGroup

	for i := 0; i < loops; i++ {
		wg.Add(1)
		go func(ctx context.Context, wg *sync.WaitGroup) {
			for j := 0; j < loops/concurrency; j++ {
				beforeLocal := time.Now()
				sqlCtx, sqlCncl := context.WithTimeout(ctx, 100*time.Millisecond)
				defer sqlCncl()
				_, errSel := status.SelectEntityoneOneByPK(sqlCtx, dbConn, benchSQLLink, testEntityoneIDs[i%len(testEntityoneIDs)])
				if errSel != nil {
					errorC <- errSel
				} else {
					latenciesC <- time.Since(beforeLocal)
				}
			}
			wg.Done()
		}(ctx, &wg)
	}

	wg.Wait()
	doneC <- true
	rawRes := <-resultsC
	timeTaken := time.Since(before)

	return BenchResult{
			Action:      "selectEntityoneOneByPK",
			Loops:       loops,
			Concurrency: concurrency,
			PauseTime:   0,
			Errors:      rawRes.errCount,
			Min:         bench.GetMin(rawRes.latencies),
			Max:         bench.GetMax(rawRes.latencies),
			Median:      bench.GetMedian(rawRes.latencies),
			StandDev:    bench.GetStandardDeviation(rawRes.latencies),
			Throughput:  int(float64(loops) / timeTaken.Seconds()),
		},
		nil
}

type rawResults struct {
	latencies []time.Duration
	errCount  int
}

// handleResults will generate two channels that will receive latencies and errors
func handleResults() (chan time.Duration, chan error, chan bool, chan rawResults) {
	latenciesC := make(chan time.Duration)
	errorC := make(chan error)
	resultsC := make(chan rawResults)
	doneC := make(chan bool)

	var latencies []time.Duration
	var errCount int
	var mux sync.Mutex

	go func() {
		for {
			select {
			case latency := <-latenciesC:
				latencies = append(latencies, latency)
			case erRrrR := <-errorC:
				fmt.Println(erRrrR)
				mux.Lock()
				errCount++
				mux.Unlock()
			case <-doneC:
				resultsC <- rawResults{latencies: latencies, errCount: errCount}
			}
		}
	}()

	return latenciesC, errorC, doneC, resultsC
}

const (
	maxPauseTime = 200 * time.Millisecond
	minPauseTime = 1 * time.Millisecond
)

// dynPauseTimeInit generates a channel that will be used to dynamically update the pause time between transactions
func dynPauseTimeInit(dynPauseTime *time.Duration) chan time.Duration {
	dynPauseTimeC := make(chan time.Duration)
	go func() {
		for additionalPauseTime := range dynPauseTimeC {
			if (*dynPauseTime+additionalPauseTime) > minPauseTime && (*dynPauseTime+additionalPauseTime) < maxPauseTime {
				*dynPauseTime += additionalPauseTime
			}
		}
	}()
	return dynPauseTimeC
}
<commit_msg>removing error log<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"sync"

	"time"

	"encoding/json"

	"github.com/jmoiron/sqlx"
	"github.com/vincentserpoul/playwithsql/bench"
	"github.com/vincentserpoul/playwithsql/dbhandler"
	"github.com/vincentserpoul/playwithsql/status"
)

// Number of retries after query returns an error
const maxRetryCount = 3

// Results to be returned
type Results struct {
	DBType       string
	MaxConns     int
	Date         time.Time
	BenchResults []BenchResult
}

// BenchResult data
type BenchResult struct {
	Action      string
	Loops       int
	Concurrency int
	PauseTime   time.Duration
	Errors      int
	Min         time.Duration
	Max         time.Duration
	Median      time.Duration
	StandDev    time.Duration
	Throughput  int
}

func main() {

	// Flags
	dbName := "playwithsql"
	dbType := flag.String("db", "mysql", "type of db to bench: mysql, cockroachdb, postgres...")
	schemaType := flag.String("sch", "islatest", "type of schema to use, is latest, latest status...")
	dbHost := flag.String("host", "127.0.0.1", "host IP")
	loops := flag.Int("loops", 100, "number of loops")
	maxConns := flag.Int("maxconns", 10, "number of max connections")
	concurrency := flag.Int("concurrency", 100, "number of go routines created")

	flag.Parse()

	if *concurrency > *loops {
		*concurrency = *loops
	}

	db, err := dbhandler.Get(*dbType, *dbHost, dbName)
	if err != nil {
		log.Fatalf("%s - %s - %s, \n%v", *dbType, *dbHost, dbName, err)
	}

	// Connection
	islatestSQLLink := status.GetSQLIntImpl(*dbType, *schemaType)

	ctx := context.Background()
	err = islatestSQLLink.MigrateDown(ctx, db)
	if err != nil {
		log.Fatalf("%v", err)
	}
	err = islatestSQLLink.MigrateUp(ctx, db)
	if err != nil {
		log.Fatalf("%v", err)
	}

	// Number of max connections
	// TODO set the param in the db config
	db.SetMaxOpenConns(*maxConns)
	db.SetMaxIdleConns(*maxConns)

	var results = Results{
		DBType:   *dbType,
		MaxConns: *maxConns,
		Date:     time.Now(),
	}

	// Create
	createResults, testEntityoneIDs, err := BenchmarkCreate(ctx, *loops, *concurrency, db, islatestSQLLink)
	if err != nil {
		log.Fatalf("%v", err)
	}
	results.BenchResults = append(results.BenchResults, createResults)

	// Update
	updateResults, err := BenchmarkUpdateStatus(ctx, *loops, *concurrency, db, islatestSQLLink, testEntityoneIDs)
	if err != nil {
		log.Fatalf("%v", err)
	}
	results.BenchResults = append(results.BenchResults, updateResults)

	// Select by status
	selectByStatusResults, err := BenchmarkSelectEntityoneByStatus(ctx, *loops, *concurrency, db, islatestSQLLink)
	if err != nil {
		log.Fatalf("%v", err)
	}
	results.BenchResults = append(results.BenchResults, selectByStatusResults)

	// Select by PK
	selectByPKResults, err := BenchmarkSelectEntityoneOneByPK(ctx, *loops, *concurrency, db, islatestSQLLink, testEntityoneIDs)
	if err != nil {
		log.Fatalf("%v", err)
	}
	results.BenchResults = append(results.BenchResults, selectByPKResults)

	jsonResults, err := json.Marshal(results)
	if err != nil {
		log.Fatalf("%v", err)
	}

	fmt.Printf("%s\n", jsonResults)
}

// BenchmarkCreate will loop a loops number of time and give the resulting time taken
func BenchmarkCreate(
	ctx context.Context,
	loops int,
	concurrency int,
	dbConn *sqlx.DB,
	benchSQLLink *status.SQLIntImpl,
) (
	results BenchResult,
	testEntityoneIDs []int64,
	err error,
) {
	entityIDsC := make(chan int64)
	latenciesC, errorC, doneC, resultsC := handleResults()

	before := time.Now()
	var wg sync.WaitGroup

	// Pause time
	dynPauseTime := 1 * time.Millisecond
	dynPauseTimeC := dynPauseTimeInit(&dynPauseTime)
	defer close(dynPauseTimeC)

	// Launch as many concurrent connections as asked
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func(routineNum int, ctx context.Context, wg *sync.WaitGroup) {
			for j := 0; j < loops/concurrency; j++ {
				time.Sleep(dynPauseTime)
				var e status.Entityone
				beforeLocal := time.Now()
				ok := false
				var errCr error
				retryCount := 0
				for retryCount < maxRetryCount && !ok {
					// Timeout
					sqlCtx, sqlCncl := context.WithTimeout(ctx, 250*time.Millisecond)
					defer sqlCncl()

					// For each error, we add some pause time
					errCr = e.Create(sqlCtx, dbConn, benchSQLLink)
					if errCr != nil {
						retryCount++
						time.Sleep(dynPauseTime)
						dynPauseTimeC <- 1 * time.Millisecond
					} else {
						ok = true
					}
				}
				if errCr != nil {
					errorC <- errCr
				} else {
					latenciesC <- time.Since(beforeLocal)
					entityIDsC <- e.ID
					// If no error, we increment down a little bit
					dynPauseTimeC <- -1 * time.Millisecond
				}
			}
			wg.Done()
		}(i, ctx, &wg)
	}

	// Receive the entityIDs
	go func() {
		for entityID := range entityIDsC {
			testEntityoneIDs = append(testEntityoneIDs, entityID)
		}
	}()

	wg.Wait()
	doneC <- true
	rawRes := <-resultsC
	timeTaken := time.Since(before)
	return BenchResult{
			Action:      "create",
			Loops:       loops,
			Concurrency: concurrency,
			PauseTime:   dynPauseTime,
			Errors:      rawRes.errCount,
			Min:         bench.GetMin(rawRes.latencies),
			Max:         bench.GetMax(rawRes.latencies),
			Median:      bench.GetMedian(rawRes.latencies),
			StandDev:    bench.GetStandardDeviation(rawRes.latencies),
			Throughput:  int(float64(loops) / timeTaken.Seconds()),
		},
		testEntityoneIDs,
		nil
}

// BenchmarkUpdateStatus benchmark for status updates (include deletes)
func BenchmarkUpdateStatus(
	ctx context.Context,
	loops int,
	concurrency int,
	dbConn *sqlx.DB,
	benchSQLLink *status.SQLIntImpl,
	testEntityoneIDs []int64,
) (
	results BenchResult,
	err error,
) {
	if len(testEntityoneIDs) == 0 {
		return results, fmt.Errorf("BenchmarkUpdateStatus: no entity created, nothing to update")
	}

	latenciesC, errorC, doneC, resultsC := handleResults()

	before := time.Now()
	var wg sync.WaitGroup

	// Pause time
	dynPauseTime := 1 * time.Millisecond
	dynPauseTimeC := dynPauseTimeInit(&dynPauseTime)
	defer close(dynPauseTimeC)

	for i := 0; i < concurrency; i++ {
		wg.Add(1)

		go func(ctx context.Context, wg *sync.WaitGroup) {
			for j := 0; j < loops/concurrency; j++ {
				time.Sleep(dynPauseTime)
				var e status.Entityone
				e.ID = testEntityoneIDs[i%len(testEntityoneIDs)]
				beforeLocal := time.Now()
				ok := false
				var errU error
				retryCount := 0
				for retryCount < maxRetryCount && !ok {
					// Timeout
					sqlCtx, sqlCncl := context.WithTimeout(ctx, 250*time.Millisecond)
					defer sqlCncl()
					errU = e.UpdateStatus(sqlCtx, dbConn, benchSQLLink, status.ActionCancel, status.StatusCancelled)
					if errU != nil {
						retryCount++
						time.Sleep(dynPauseTime)
						dynPauseTimeC <- 1 * time.Millisecond
					} else {
						ok = true
					}
				}
				if errU != nil {
					errorC <- errU
				} else {
					latenciesC <- time.Since(beforeLocal)
					// If no error, we increment down a little bit
					dynPauseTimeC <- -1 * time.Millisecond
				}
			}
			wg.Done()
		}(ctx, &wg)
	}

	wg.Wait()
	doneC <- true
	rawRes := <-resultsC
	timeTaken := time.Since(before)

	return BenchResult{
			Action:      "updateStatus",
			Loops:       loops,
			Concurrency: concurrency,
			PauseTime:   dynPauseTime,
			Errors:      rawRes.errCount,
			Min:         bench.GetMin(rawRes.latencies),
			Max:         bench.GetMax(rawRes.latencies),
			Median:      bench.GetMedian(rawRes.latencies),
			StandDev:    bench.GetStandardDeviation(rawRes.latencies),
			Throughput:  int(float64(loops) / timeTaken.Seconds()),
		},
		nil

}

// BenchmarkSelectEntityoneByStatus benchmark with select by status
func BenchmarkSelectEntityoneByStatus(
	ctx context.Context,
	loops int,
	concurrency int,
	dbConn *sqlx.DB,
	benchSQLLink *status.SQLIntImpl,
) (
	results BenchResult,
	err error,
) {
	latenciesC, errorC, doneC, resultsC := handleResults()

	var wg sync.WaitGroup
	before := time.Now()

	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func(ctx context.Context, wg *sync.WaitGroup) {
			for j := 0; j < loops/concurrency; j++ {
				beforeLocal := time.Now()
				sqlCtx, sqlCncl := context.WithTimeout(ctx, 100*time.Millisecond)
				defer sqlCncl()
				_, errSel := status.SelectEntityoneByStatus(sqlCtx, dbConn, benchSQLLink, status.StatusCancelled)
				if errSel != nil {
					errorC <- errSel
				} else {
					latenciesC <- time.Since(beforeLocal)
				}
			}
			wg.Done()
		}(ctx, &wg)
	}

	wg.Wait()
	doneC <- true
	rawRes := <-resultsC
	timeTaken := time.Since(before)

	return BenchResult{
			Action:      "selectEntityoneByStatus",
			Loops:       loops,
			Concurrency: concurrency,
			PauseTime:   0,
			Errors:      rawRes.errCount,
			Min:         bench.GetMin(rawRes.latencies),
			Max:         bench.GetMax(rawRes.latencies),
			Median:      bench.GetMedian(rawRes.latencies),
			StandDev:    bench.GetStandardDeviation(rawRes.latencies),
			Throughput:  int(float64(loops) / timeTaken.Seconds()),
		},
		nil
}

// BenchmarkSelectEntityoneOneByPK benchmark with select by primary key
func BenchmarkSelectEntityoneOneByPK(
	ctx context.Context,
	loops int,
	concurrency int,
	dbConn *sqlx.DB,
	benchSQLLink *status.SQLIntImpl,
	testEntityoneIDs []int64,
) (
	results BenchResult,
	err error,
) {
	latenciesC, errorC, doneC, resultsC := handleResults()

	before := time.Now()
	var wg sync.WaitGroup

	for i := 0; i < loops; i++ {
		wg.Add(1)
		go func(ctx context.Context, wg *sync.WaitGroup) {
			for j := 0; j < loops/concurrency; j++ {
				beforeLocal := time.Now()
				sqlCtx, sqlCncl := context.WithTimeout(ctx, 100*time.Millisecond)
				defer sqlCncl()
				_, errSel := status.SelectEntityoneOneByPK(sqlCtx, dbConn, benchSQLLink, testEntityoneIDs[i%len(testEntityoneIDs)])
				if errSel != nil {
					errorC <- errSel
				} else {
					latenciesC <- time.Since(beforeLocal)
				}
			}
			wg.Done()
		}(ctx, &wg)
	}

	wg.Wait()
	doneC <- true
	rawRes := <-resultsC
	timeTaken := time.Since(before)

	return BenchResult{
			Action:      "selectEntityoneOneByPK",
			Loops:       loops,
			Concurrency: concurrency,
			PauseTime:   0,
			Errors:      rawRes.errCount,
			Min:         bench.GetMin(rawRes.latencies),
			Max:         bench.GetMax(rawRes.latencies),
			Median:      bench.GetMedian(rawRes.latencies),
			StandDev:    bench.GetStandardDeviation(rawRes.latencies),
			Throughput:  int(float64(loops) / timeTaken.Seconds()),
		},
		nil
}

type rawResults struct {
	latencies []time.Duration
	errCount  int
}

// handleResults will generate two channels that will receive latencies and errors
func handleResults() (chan time.Duration, chan error, chan bool, chan rawResults) {
	latenciesC := make(chan time.Duration)
	errorC := make(chan error)
	resultsC := make(chan rawResults)
	doneC := make(chan bool)

	var latencies []time.Duration
	var errCount int
	var mux sync.Mutex

	go func() {
		for {
			select {
			case latency := <-latenciesC:
				latencies = append(latencies, latency)
			// case erRrrR := <-errorC:
			// 	fmt.Println(erRrrR)
			case <-errorC:
				mux.Lock()
				errCount++
				mux.Unlock()
			case <-doneC:
				resultsC <- rawResults{latencies: latencies, errCount: errCount}
			}
		}
	}()

	return latenciesC, errorC, doneC, resultsC
}

const (
	maxPauseTime = 200 * time.Millisecond
	minPauseTime = 1 * time.Millisecond
)

// dynPauseTimeInit generates a channel that will be used to dynamically update the pause time between transactions
func dynPauseTimeInit(dynPauseTime *time.Duration) chan time.Duration {
	dynPauseTimeC := make(chan time.Duration)
	go func() {
		for additionalPauseTime := range dynPauseTimeC {
			if (*dynPauseTime+additionalPauseTime) > minPauseTime && (*dynPauseTime+additionalPauseTime) < maxPauseTime {
				*dynPauseTime += additionalPauseTime
			}
		}
	}()
	return dynPauseTimeC
}
<|endoftext|>


<commit_before>package builder

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"sort"
)

func (b *Builder) writeHtml(writer io.Writer) {
	fmt.Fprintf(writer,
		"Number of image streams: %d
\n",
		b.getNumNormalStreams())
	currentBuilds := make([]string, 0)
	goodBuilds := make(map[string]buildResultType)
	failedBuilds := make(map[string]buildResultType)
	b.buildResultsLock.RLock()
	for name := range b.currentBuildLogs {
		currentBuilds = append(currentBuilds, name)
	}
	for name, result := range b.lastBuildResults {
		if result.error == nil {
			goodBuilds[name] = result
		} else {
			failedBuilds[name] = result
		}
	}
	b.buildResultsLock.RUnlock()
	if len(currentBuilds) > 0 {
		fmt.Fprintln(writer, "Current image builds:
")
		fmt.Fprintln(writer, `<table border="1">`)
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "    Image Stream	")
		fmt.Fprintln(writer, "    Build log	")
		fmt.Fprintln(writer, "
")
		for _, streamName := range currentBuilds {
			fmt.Fprintf(writer, "
\n")
			fmt.Fprintf(writer, "    %s	\n", streamName)
			fmt.Fprintf(writer,
				"    log	\n",
				streamName)
			fmt.Fprintf(writer, "
\n")
		}
		fmt.Fprintln(writer, "

")
	}
	if len(failedBuilds) > 0 {
		streamNames := make([]string, 0, len(failedBuilds))
		for streamName := range failedBuilds {
			streamNames = append(streamNames, streamName)
		}
		sort.Strings(streamNames)
		fmt.Fprintln(writer, "Failed image builds:
")
		fmt.Fprintln(writer, `<table border="1">`)
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "    Image Stream	")
		fmt.Fprintln(writer, "    Error	")
		fmt.Fprintln(writer, "    Build log	")
		fmt.Fprintln(writer, "
")
		for _, streamName := range streamNames {
			result := failedBuilds[streamName]
			fmt.Fprintf(writer, "
\n")
			fmt.Fprintf(writer, "    %s	\n", streamName)
			fmt.Fprintf(writer, "    %s	\n", result.error)
			fmt.Fprintf(writer,
				"    log	\n",
				streamName)
			fmt.Fprintf(writer, "
\n")
		}
		fmt.Fprintln(writer, "

")
	}
	if len(goodBuilds) > 0 {
		streamNames := make([]string, 0, len(goodBuilds))
		for streamName := range goodBuilds {
			streamNames = append(streamNames, streamName)
		}
		sort.Strings(streamNames)
		fmt.Fprintln(writer, "Successful image builds:
")
		fmt.Fprintln(writer, `<table border="1">`)
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "    Image Stream	")
		fmt.Fprintln(writer, "    Name	")
		fmt.Fprintln(writer, "    Build log	")
		fmt.Fprintln(writer, "
")
		for _, streamName := range streamNames {
			result := goodBuilds[streamName]
			fmt.Fprintf(writer, "
\n")
			fmt.Fprintf(writer, "    %s	\n", streamName)
			fmt.Fprintf(writer,
				"    %s	\n",
				b.imageServerAddress, result.imageName, result.imageName)
			fmt.Fprintf(writer,
				"    log	\n",
				streamName)
			fmt.Fprintf(writer, "
\n")
		}
		fmt.Fprintln(writer, "

")
	}
}

func (b *Builder) showImageStream(writer io.Writer, streamName string) {
	stream := b.getNormalStream(streamName)
	if stream == nil {
		fmt.Fprintf(writer, "Stream: %s does not exist!\n", streamName)
		return
	}
	fmt.Fprintf(writer, "
Information for stream: %s
\n", streamName)
	fmt.Fprintf(writer, "Manifest URL: %s
\n",
		stream.ManifestUrl)
	fmt.Fprintf(writer, "Manifest Directory: %s
\n",
		stream.ManifestDirectory)
	manifestRoot, manifestDirectory, err := stream.getManifest(b, streamName,
		"", new(bytes.Buffer))
	if err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	}
	defer os.RemoveAll(manifestRoot)
	manifestFilename := path.Join(manifestRoot, manifestDirectory, "manifest")
	manifestBytes, err := ioutil.ReadFile(manifestFilename)
	if err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	}
	var manifest manifestType
	if err := json.Unmarshal(manifestBytes, &manifest); err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	}
	sourceStream := b.getNormalStream(manifest.SourceImage)
	if sourceStream == nil {
		fmt.Fprintf(writer, "SourceImage: %s
\n",
			manifest.SourceImage)
	} else {
		fmt.Fprintf(writer,
			"SourceImage: %s
\n",
			manifest.SourceImage, manifest.SourceImage)
	}
	fmt.Fprintln(writer, "Contents of manifest file:
")
	fmt.Fprintln(writer, "
")
	writer.Write(manifestBytes)
	fmt.Fprintln(writer, "

")
	packagesFile, err := os.Open(
		path.Join(manifestRoot, manifestDirectory, "package-list"))
	if err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	}
	defer packagesFile.Close()
	fmt.Fprintln(writer, "Contents of package-list file:
")
	fmt.Fprintln(writer, "
")
	io.Copy(writer, packagesFile)
	fmt.Fprintln(writer, "

")
}

func (b *Builder) showImageStreams(writer io.Writer) {
	streamNames := b.listNormalStreamNames()
	sort.Strings(streamNames)
	fmt.Fprintln(writer, `<table border="1">`)
	fmt.Fprintln(writer, "
")
	fmt.Fprintln(writer, "    Image Stream	")
	fmt.Fprintln(writer, "    ManifestUrl	")
	fmt.Fprintln(writer, "    ManifestDirectory	")
	fmt.Fprintln(writer, "
")
	for _, streamName := range streamNames {
		imageStream := b.getNormalStream(streamName)
		if imageStream == nil {
			continue
		}
		fmt.Fprintf(writer, "
\n")
		fmt.Fprintf(writer,
			"    %s	\n",
			streamName, streamName)
		fmt.Fprintf(writer, "    %s	\n", imageStream.ManifestUrl)
		fmt.Fprintf(writer, "    %s	\n", imageStream.ManifestDirectory)
		fmt.Fprintf(writer, "
\n")
	}
	fmt.Fprintln(writer, "

")
}
<commit_msg>Show repository and manifest sizes when showing image stream in imaginator.<commit_after>package builder

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"sort"
	"time"

	"github.com/Symantec/Dominator/lib/format"
)

func (b *Builder) writeHtml(writer io.Writer) {
	fmt.Fprintf(writer,
		"Number of image streams: %d
\n",
		b.getNumNormalStreams())
	currentBuilds := make([]string, 0)
	goodBuilds := make(map[string]buildResultType)
	failedBuilds := make(map[string]buildResultType)
	b.buildResultsLock.RLock()
	for name := range b.currentBuildLogs {
		currentBuilds = append(currentBuilds, name)
	}
	for name, result := range b.lastBuildResults {
		if result.error == nil {
			goodBuilds[name] = result
		} else {
			failedBuilds[name] = result
		}
	}
	b.buildResultsLock.RUnlock()
	if len(currentBuilds) > 0 {
		fmt.Fprintln(writer, "Current image builds:
")
		fmt.Fprintln(writer, `<table border="1">`)
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "    Image Stream	")
		fmt.Fprintln(writer, "    Build log	")
		fmt.Fprintln(writer, "
")
		for _, streamName := range currentBuilds {
			fmt.Fprintf(writer, "
\n")
			fmt.Fprintf(writer, "    %s	\n", streamName)
			fmt.Fprintf(writer,
				"    log	\n",
				streamName)
			fmt.Fprintf(writer, "
\n")
		}
		fmt.Fprintln(writer, "

")
	}
	if len(failedBuilds) > 0 {
		streamNames := make([]string, 0, len(failedBuilds))
		for streamName := range failedBuilds {
			streamNames = append(streamNames, streamName)
		}
		sort.Strings(streamNames)
		fmt.Fprintln(writer, "Failed image builds:
")
		fmt.Fprintln(writer, `<table border="1">`)
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "    Image Stream	")
		fmt.Fprintln(writer, "    Error	")
		fmt.Fprintln(writer, "    Build log	")
		fmt.Fprintln(writer, "
")
		for _, streamName := range streamNames {
			result := failedBuilds[streamName]
			fmt.Fprintf(writer, "
\n")
			fmt.Fprintf(writer, "    %s	\n", streamName)
			fmt.Fprintf(writer, "    %s	\n", result.error)
			fmt.Fprintf(writer,
				"    log	\n",
				streamName)
			fmt.Fprintf(writer, "
\n")
		}
		fmt.Fprintln(writer, "

")
	}
	if len(goodBuilds) > 0 {
		streamNames := make([]string, 0, len(goodBuilds))
		for streamName := range goodBuilds {
			streamNames = append(streamNames, streamName)
		}
		sort.Strings(streamNames)
		fmt.Fprintln(writer, "Successful image builds:
")
		fmt.Fprintln(writer, `<table border="1">`)
		fmt.Fprintln(writer, "
")
		fmt.Fprintln(writer, "    Image Stream	")
		fmt.Fprintln(writer, "    Name	")
		fmt.Fprintln(writer, "    Build log	")
		fmt.Fprintln(writer, "
")
		for _, streamName := range streamNames {
			result := goodBuilds[streamName]
			fmt.Fprintf(writer, "
\n")
			fmt.Fprintf(writer, "    %s	\n", streamName)
			fmt.Fprintf(writer,
				"    %s	\n",
				b.imageServerAddress, result.imageName, result.imageName)
			fmt.Fprintf(writer,
				"    log	\n",
				streamName)
			fmt.Fprintf(writer, "
\n")
		}
		fmt.Fprintln(writer, "

")
	}
}

func (b *Builder) showImageStream(writer io.Writer, streamName string) {
	stream := b.getNormalStream(streamName)
	if stream == nil {
		fmt.Fprintf(writer, "Stream: %s does not exist!\n", streamName)
		return
	}
	fmt.Fprintf(writer, "
Information for stream: %s
\n", streamName)
	fmt.Fprintf(writer, "Manifest URL: %s
\n",
		stream.ManifestUrl)
	fmt.Fprintf(writer, "Manifest Directory: %s
\n",
		stream.ManifestDirectory)
	startTime := time.Now()
	manifestRoot, manifestDirectory, err := stream.getManifest(b, streamName,
		"", new(bytes.Buffer))
	if err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	}
	defer os.RemoveAll(manifestRoot)
	loadTime := time.Since(startTime)
	manifestDirectory = path.Join(manifestRoot, manifestDirectory)
	manifestFilename := path.Join(manifestDirectory, "manifest")
	manifestBytes, err := ioutil.ReadFile(manifestFilename)
	if err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	}
	var manifest manifestType
	if err := json.Unmarshal(manifestBytes, &manifest); err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	}
	sourceStream := b.getNormalStream(manifest.SourceImage)
	if sourceStream == nil {
		fmt.Fprintf(writer, "SourceImage: %s
\n",
			manifest.SourceImage)
	} else {
		fmt.Fprintf(writer,
			"SourceImage: %s
\n",
			manifest.SourceImage, manifest.SourceImage)
	}
	fmt.Fprintln(writer, "Contents of manifest file:
")
	fmt.Fprintln(writer, "
")
	writer.Write(manifestBytes)
	fmt.Fprintln(writer, "

")
	packagesFile, err := os.Open(path.Join(manifestDirectory, "package-list"))
	if err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	}
	defer packagesFile.Close()
	fmt.Fprintln(writer, "Contents of package-list file:
")
	fmt.Fprintln(writer, "
")
	io.Copy(writer, packagesFile)
	fmt.Fprintln(writer, "

")
	if size, err := getTreeSize(manifestRoot); err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	} else {
		speed := float64(size) / loadTime.Seconds()
		fmt.Fprintf(writer, "Repository size: %s, took: %s (%s/s)
\n",
			format.FormatBytes(size), format.Duration(loadTime),
			format.FormatBytes(uint64(speed)))
	}
	if size, err := getTreeSize(manifestDirectory); err != nil {
		fmt.Fprintf(writer, "%s
\n", err)
		return
	} else {
		fmt.Fprintf(writer, "Manifest tree size: %s
\n",
			format.FormatBytes(size))
	}
}

func (b *Builder) showImageStreams(writer io.Writer) {
	streamNames := b.listNormalStreamNames()
	sort.Strings(streamNames)
	fmt.Fprintln(writer, `<table border="1">`)
	fmt.Fprintln(writer, "
")
	fmt.Fprintln(writer, "    Image Stream	")
	fmt.Fprintln(writer, "    ManifestUrl	")
	fmt.Fprintln(writer, "    ManifestDirectory	")
	fmt.Fprintln(writer, "
")
	for _, streamName := range streamNames {
		imageStream := b.getNormalStream(streamName)
		if imageStream == nil {
			continue
		}
		fmt.Fprintf(writer, "
\n")
		fmt.Fprintf(writer,
			"    %s	\n",
			streamName, streamName)
		fmt.Fprintf(writer, "    %s	\n", imageStream.ManifestUrl)
		fmt.Fprintf(writer, "    %s	\n", imageStream.ManifestDirectory)
		fmt.Fprintf(writer, "
\n")
	}
	fmt.Fprintln(writer, "

")
}

func getTreeSize(dirname string) (uint64, error) {
	var size uint64
	err := filepath.Walk(dirname,
		func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			size += uint64(info.Size())
			return nil
		})
	if err != nil {
		return 0, err
	}
	return size, nil
}
<|endoftext|>


<commit_before>package socks

import (
	"io"
	"net"
	"strconv"
	"sync"
	"time"

	"golang.org/x/net/proxy"

	"github.com/eycorsican/go-tun2socks/common/dns"
	"github.com/eycorsican/go-tun2socks/common/log"
	"github.com/eycorsican/go-tun2socks/common/lsof"
	"github.com/eycorsican/go-tun2socks/common/stats"
	"github.com/eycorsican/go-tun2socks/core"
)

type tcpHandler struct {
	sync.Mutex

	proxyHost string
	proxyPort uint16

	fakeDns       dns.FakeDns
	sessionStater stats.SessionStater
}

func NewTCPHandler(proxyHost string, proxyPort uint16, fakeDns dns.FakeDns, sessionStater stats.SessionStater) core.TCPConnHandler {
	return &tcpHandler{
		proxyHost:     proxyHost,
		proxyPort:     proxyPort,
		fakeDns:       fakeDns,
		sessionStater: sessionStater,
	}
}

type direction byte

const (
	dirUplink direction = iota
	dirDownlink
)

func statsCopy(dst io.Writer, src io.Reader, sess *stats.Session, dir direction) (written int64, err error) {
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				switch dir {
				case dirUplink:
					sess.AddUploadBytes(int64(nw))
				case dirDownlink:
					sess.AddDownloadBytes(int64(nw))
				default:
				}
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er != nil {
			if er != io.EOF {
				err = er
			}
			break
		}
	}
	return written, err
}

type duplexConn interface {
	net.Conn
	CloseRead() error
	CloseWrite() error
}

func (h *tcpHandler) relay(lhs, rhs net.Conn, sess *stats.Session) {
	var err error
	upCh := make(chan struct{})

	cls := func(dir direction) {
		lhsDConn, lhsOk := lhs.(duplexConn)
		rhsDConn, rhsOk := rhs.(duplexConn)
		if lhsOk && rhsOk {
			switch dir {
			case dirUplink:
				lhsDConn.CloseRead()
				rhsDConn.CloseWrite()
			case dirDownlink:
				lhsDConn.CloseWrite()
				rhsDConn.CloseRead()
			default:
				panic("unexpected direction")
			}
		} else {
			lhs.Close()
			rhs.Close()
		}
	}

	go func() {
		if h.sessionStater != nil && sess != nil {
			_, err = statsCopy(rhs, lhs, sess, dirUplink)
		} else {
			_, err = io.Copy(rhs, lhs)
		}
		if err != nil && err != io.EOF {
			cls(dirUplink)
		}
		upCh <- struct{}{}
	}()

	if h.sessionStater != nil && sess != nil {
		_, err = statsCopy(lhs, rhs, sess, dirDownlink)
	} else {
		_, err = io.Copy(rhs, lhs)
	}
	if err != nil && err != io.EOF {
		cls(dirDownlink)
	}

	<-upCh // Wait for uplink done.

	if h.sessionStater != nil {
		h.sessionStater.RemoveSession(lhs)
	}
}

func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
	dialer, err := proxy.SOCKS5("tcp", core.ParseTCPAddr(h.proxyHost, h.proxyPort).String(), nil, nil)
	if err != nil {
		return err
	}

	// Replace with a domain name if target address IP is a fake IP.
	var targetHost string
	if h.fakeDns != nil && h.fakeDns.IsFakeIP(target.IP) {
		targetHost = h.fakeDns.QueryDomain(target.IP)
	} else {
		targetHost = target.IP.String()
	}
	dest := net.JoinHostPort(targetHost, strconv.Itoa(target.Port))

	c, err := dialer.Dial(target.Network(), dest)
	if err != nil {
		return err
	}

	var process string
	var sess *stats.Session
	if h.sessionStater != nil {
		// Get name of the process.
		localHost, localPortStr, _ := net.SplitHostPort(conn.LocalAddr().String())
		localPortInt, _ := strconv.Atoi(localPortStr)
		process, err = lsof.GetCommandNameBySocket(target.Network(), localHost, uint16(localPortInt))
		if err != nil {
			process = "unknown process"
		}

		sess = &stats.Session{
			process,
			target.Network(),
			conn.LocalAddr().String(),
			dest,
			0,
			0,
			time.Now(),
		}
		h.sessionStater.AddSession(conn, sess)
	}

	go h.relay(conn, c, sess)

	log.Access(process, "proxy", target.Network(), conn.LocalAddr().String(), dest)

	return nil
}
<commit_msg>interrupt the connection if error is not EOF<commit_after>package socks

import (
	"io"
	"net"
	"strconv"
	"sync"
	"time"

	"golang.org/x/net/proxy"

	"github.com/eycorsican/go-tun2socks/common/dns"
	"github.com/eycorsican/go-tun2socks/common/log"
	"github.com/eycorsican/go-tun2socks/common/lsof"
	"github.com/eycorsican/go-tun2socks/common/stats"
	"github.com/eycorsican/go-tun2socks/core"
)

type tcpHandler struct {
	sync.Mutex

	proxyHost string
	proxyPort uint16

	fakeDns       dns.FakeDns
	sessionStater stats.SessionStater
}

func NewTCPHandler(proxyHost string, proxyPort uint16, fakeDns dns.FakeDns, sessionStater stats.SessionStater) core.TCPConnHandler {
	return &tcpHandler{
		proxyHost:     proxyHost,
		proxyPort:     proxyPort,
		fakeDns:       fakeDns,
		sessionStater: sessionStater,
	}
}

type direction byte

const (
	dirUplink direction = iota
	dirDownlink
)

func statsCopy(dst io.Writer, src io.Reader, sess *stats.Session, dir direction) (written int64, err error) {
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				switch dir {
				case dirUplink:
					sess.AddUploadBytes(int64(nw))
				case dirDownlink:
					sess.AddDownloadBytes(int64(nw))
				default:
				}
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er != nil {
			if er != io.EOF {
				err = er
			}
			break
		}
	}
	return written, err
}

type duplexConn interface {
	net.Conn
	CloseRead() error
	CloseWrite() error
}

func (h *tcpHandler) relay(lhs, rhs net.Conn, sess *stats.Session) {
	var err error
	upCh := make(chan struct{})

	cls := func(dir direction, interrupt bool) {
		lhsDConn, lhsOk := lhs.(duplexConn)
		rhsDConn, rhsOk := rhs.(duplexConn)
		if !interrupt && lhsOk && rhsOk {
			switch dir {
			case dirUplink:
				lhsDConn.CloseRead()
				rhsDConn.CloseWrite()
			case dirDownlink:
				lhsDConn.CloseWrite()
				rhsDConn.CloseRead()
			default:
				panic("unexpected direction")
			}
		} else {
			lhs.Close()
			rhs.Close()
		}
	}

	// Uplink
	go func() {
		if h.sessionStater != nil && sess != nil {
			_, err = statsCopy(rhs, lhs, sess, dirUplink)
		} else {
			_, err = io.Copy(rhs, lhs)
		}
		if err != nil && err != io.EOF {
			cls(dirUplink, true) // interrupt the conn if the error is not EOF
		} else {
			cls(dirUplink, false) // half close uplink direction of the TCP conn if possible
		}
		upCh <- struct{}{}
	}()

	// Downlink
	if h.sessionStater != nil && sess != nil {
		_, err = statsCopy(lhs, rhs, sess, dirDownlink)
	} else {
		_, err = io.Copy(rhs, lhs)
	}
	if err != nil && err != io.EOF {
		cls(dirDownlink, true)
	} else {
		cls(dirDownlink, false)
	}

	<-upCh // Wait for uplink done.

	if h.sessionStater != nil {
		h.sessionStater.RemoveSession(lhs)
	}
}

func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
	dialer, err := proxy.SOCKS5("tcp", core.ParseTCPAddr(h.proxyHost, h.proxyPort).String(), nil, nil)
	if err != nil {
		return err
	}

	// Replace with a domain name if target address IP is a fake IP.
	var targetHost string
	if h.fakeDns != nil && h.fakeDns.IsFakeIP(target.IP) {
		targetHost = h.fakeDns.QueryDomain(target.IP)
	} else {
		targetHost = target.IP.String()
	}
	dest := net.JoinHostPort(targetHost, strconv.Itoa(target.Port))

	c, err := dialer.Dial(target.Network(), dest)
	if err != nil {
		return err
	}

	var process string
	var sess *stats.Session
	if h.sessionStater != nil {
		// Get name of the process.
		localHost, localPortStr, _ := net.SplitHostPort(conn.LocalAddr().String())
		localPortInt, _ := strconv.Atoi(localPortStr)
		process, err = lsof.GetCommandNameBySocket(target.Network(), localHost, uint16(localPortInt))
		if err != nil {
			process = "unknown process"
		}

		sess = &stats.Session{
			process,
			target.Network(),
			conn.LocalAddr().String(),
			dest,
			0,
			0,
			time.Now(),
		}
		h.sessionStater.AddSession(conn, sess)
	}

	go h.relay(conn, c, sess)

	log.Access(process, "proxy", target.Network(), conn.LocalAddr().String(), dest)

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha1"
	"fmt"
	"html/template"
	"net/http"
)

var inputTemplate = "html/input.html"

func inputHandler(w http.ResponseWriter, r *http.Request) {
	t, _ := template.ParseFiles(inputTemplate)
	t.Execute(w, nil)
}

func solveHandler(w http.ResponseWriter, r *http.Request) {
	model := r.FormValue("model")
	hash := fmt.Sprintf("%x", sha1.Sum([]byte(model)))
	// TODO: solve problem

	http.Redirect(w, r, "/result/"+hash, http.StatusFound)
}

func main() {
	http.HandleFunc("/", inputHandler)
	http.ListenAndServe(":8080", nil)
	http.HandleFunc("/solve", solveHandler)
}
<commit_msg>add log.Fatal in main<commit_after>package main

import (
	"crypto/sha1"
	"fmt"
	"html/template"
	"log"
	"net/http"
)

var inputTemplate = "html/input.html"

func inputHandler(w http.ResponseWriter, r *http.Request) {
	t, _ := template.ParseFiles(inputTemplate)
	t.Execute(w, nil)
}

func solveHandler(w http.ResponseWriter, r *http.Request) {
	model := r.FormValue("model")
	hash := fmt.Sprintf("%x", sha1.Sum([]byte(model)))
	// TODO: solve problem

	http.Redirect(w, r, "/result/"+hash, http.StatusFound)
}

func main() {
	http.HandleFunc("/", inputHandler)
	http.HandleFunc("/solve", solveHandler)
	log.Fatal(http.ListenAndServe(":8080", nil))
}
<|endoftext|>


<commit_before>package node

import (
	"bytes"
	"context"
	"encoding/gob"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"gitlab.fg/go/disco/multicast"
)

// Modes of a node being registered
const (
	RegisterAction   = iota
	DeregisterAction = iota
)

// Node represents a machine registered with Disco
type Node struct {
	Values       Values
	SrcIP        net.IP
	SendInterval time.Duration
	Action       int
	ipv6         net.IP // set by localIPv4 function
	ipv4         net.IP // set by localIPv6 function
	mc           *multicast.Multicast
	mu           *sync.Mutex // protect ipv4, ipv6, mc, SendInterval, registerCh
	registerCh   chan struct{}
}

// Values stores any values passed to the node
type Values map[string]string

func (n *Node) init() {
	if n.mu == nil {
		n.mu = &sync.Mutex{}
	}
	n.mu.Lock()
	defer n.mu.Unlock()
	if n.registerCh == nil {
		n.registerCh = make(chan struct{})
	}
}

func (n *Node) String() string {
	return fmt.Sprintf("IPv4: %s, IPv6: %s, Values: %s", n.ipv4, n.ipv6, n.Values)
}

// Equal compares nodes
func (n *Node) Equal(b *Node) bool {
	if n.mu == nil {
		n.mu = &sync.Mutex{}
	}
	if b.mu == nil {
		b.mu = &sync.Mutex{}
	}

	n.mu.Lock()
	b.mu.Lock()
	defer n.mu.Unlock()
	defer b.mu.Unlock()

	if !n.ipv4.Equal(b.ipv4) {
		return false
	}
	if !n.ipv6.Equal(b.ipv6) {
		return false
	}

	// Check if the Values map is the same
	if len(n.Values) != len(b.Values) {
		return false
	}
	for k := range n.Values {
		v1 := n.Values[k]
		v2 := b.Values[k]
		if v1 != v2 {
			return false
		}
	}

	return true
}

// GobEncode gob interface
func (n *Node) GobEncode() ([]byte, error) {
	if n.mu == nil {
		n.mu = &sync.Mutex{}
	}

	n.mu.Lock()
	defer n.mu.Unlock()

	w := new(bytes.Buffer)
	encoder := gob.NewEncoder(w)

	if err := encoder.Encode(n.SendInterval); err != nil {
		return nil, err
	}

	if err := encoder.Encode(n.ipv4); err != nil {
		return nil, err
	}

	if err := encoder.Encode(n.ipv6); err != nil {
		return nil, err
	}

	if err := encoder.Encode(n.Values); err != nil {
		return nil, err
	}

	return w.Bytes(), nil
}

// GobDecode gob interface
func (n *Node) GobDecode(buf []byte) error {
	if n.mu == nil {
		n.mu = &sync.Mutex{}
	}

	n.mu.Lock()
	defer n.mu.Unlock()

	r := bytes.NewBuffer(buf)
	decoder := gob.NewDecoder(r)

	if err := decoder.Decode(&n.SendInterval); err != nil {
		return err
	}

	if err := decoder.Decode(&n.ipv4); err != nil {
		return err
	}

	if err := decoder.Decode(&n.ipv6); err != nil {
		return err
	}

	return decoder.Decode(&n.Values)
}

// Done returns a channel that can be used to wait till Multicast is stopped
func (n *Node) Done() <-chan struct{} {
	return n.mc.Done()
}

// RegisterCh returns a channel to know if the node should stay registered
func (n *Node) RegisterCh() <-chan struct{} {
	n.init()
	return n.registerCh
}

// KeepRegistered sends an anonymous struct{} to registeredChan to indicate the node should stay registered
func (n *Node) KeepRegistered() {
	n.init()
	n.registerCh <- struct{}{}
}

// Multicast start the mulicast ping
func (n *Node) Multicast(ctx context.Context, multicastAddress string) error {
	if n.mu == nil {
		n.mu = &sync.Mutex{}
	}
	n.mu.Lock()
	n.ipv4 = localIPv4()
	n.ipv6 = localIPv6()

	if n.SendInterval.Seconds() == float64(0) {
		n.SendInterval = 1 * time.Second // default to 1 second
	}

	n.mu.Unlock()

	// Encode node to be sent via multicast
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(n)
	if err != nil {
		return err
	}

	n.mu.Lock()
	defer n.mu.Unlock()
	n.mc = &multicast.Multicast{Address: multicastAddress}
	if err := n.mc.Send(ctx, n.SendInterval, buf.Bytes()); err != nil {
		return err
	}

	return nil
}

// Stop closes the StopCh to stop multicast sending
func (n *Node) Stop() {
	if n.mu == nil {
		n.mu = &sync.Mutex{}
	}
	n.mu.Lock()
	defer n.mu.Unlock()
	n.mc.Stop()
}

// IPv4 getter for ipv4Address
func (n *Node) IPv4() net.IP {
	if n.mu == nil {
		n.mu = &sync.Mutex{}
	}
	n.mu.Lock()
	defer n.mu.Unlock()
	return n.ipv4
}

// IPv6 getter for ipv6Address
func (n *Node) IPv6() net.IP {
	if n.mu == nil {
		n.mu = &sync.Mutex{}
	}
	n.mu.Lock()
	defer n.mu.Unlock()
	return n.ipv6
}

// localIPv4 return the ipv4 address of the computer
// If it can't get the local ip it returns 127.0.0.1
// https://github.com/forestgiant/netutil
func localIPv4() net.IP {
	loopback := net.ParseIP("127.0.0.1")

	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return loopback
	}

	for _, addr := range addrs {
		// check the address type and make sure it's not loopback
		if ipnet, ok := addr.(*net.IPNet); ok {
			if !ipnet.IP.IsLoopback() {
				if ipnet.IP.To4() != nil {
					return ipnet.IP.To4()
				}
			}
		}
	}

	return loopback
}

// localIPv6 return the ipv6 address of the computer
// If it can't get the local ip it returns net.IPv6loopback
// https://github.com/forestgiant/netutil
func localIPv6() net.IP {
	loopback := net.IPv6loopback

	intfs, err := net.Interfaces()
	if err != nil {
		return loopback
	}

	for _, intf := range intfs {
		// If the interface is a loopback or doesn't have multicasting let's skip it
		if strings.Contains(intf.Flags.String(), net.FlagLoopback.String()) || !strings.Contains(intf.Flags.String(), net.FlagMulticast.String()) {
			continue
		}

		// Now let's check if the interface has an ipv6 address
		addrs, err := intf.Addrs()
		if err != nil {
			continue
		}

		for _, address := range addrs {
			if ipnet, ok := address.(*net.IPNet); ok {
				if !ipnet.IP.IsLoopback() {
					if ipnet.IP.To4() == nil {
						return ipnet.IP
					}
				}
			}
		}
	}

	return loopback
}
<commit_msg>Fixing race condition in node.go because of checking if a mutex was nil<commit_after>package node

import (
	"bytes"
	"context"
	"encoding/gob"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"gitlab.fg/go/disco/multicast"
)

// Modes of a node being registered
const (
	RegisterAction   = iota
	DeregisterAction = iota
)

// Node represents a machine registered with Disco
type Node struct {
	Values       Values
	SrcIP        net.IP
	SendInterval time.Duration
	Action       int
	ipv6         net.IP // set by localIPv4 function
	ipv4         net.IP // set by localIPv6 function
	mc           *multicast.Multicast
	mu           sync.Mutex // protect ipv4, ipv6, mc, SendInterval, registerCh
	registerCh   chan struct{}
}

// Values stores any values passed to the node
type Values map[string]string

func (n *Node) init() {
	n.mu.Lock()
	defer n.mu.Unlock()
	if n.registerCh == nil {
		n.registerCh = make(chan struct{})
	}
}

func (n *Node) String() string {
	return fmt.Sprintf("IPv4: %s, IPv6: %s, Values: %s", n.ipv4, n.ipv6, n.Values)
}

// Equal compares nodes
func (n *Node) Equal(b *Node) bool {
	n.mu.Lock()
	b.mu.Lock()
	defer n.mu.Unlock()
	defer b.mu.Unlock()

	if !n.ipv4.Equal(b.ipv4) {
		return false
	}
	if !n.ipv6.Equal(b.ipv6) {
		return false
	}

	// Check if the Values map is the same
	if len(n.Values) != len(b.Values) {
		return false
	}
	for k := range n.Values {
		v1 := n.Values[k]
		v2 := b.Values[k]
		if v1 != v2 {
			return false
		}
	}

	return true
}

// GobEncode gob interface
func (n *Node) GobEncode() ([]byte, error) {
	n.mu.Lock()
	defer n.mu.Unlock()

	w := new(bytes.Buffer)
	encoder := gob.NewEncoder(w)

	if err := encoder.Encode(n.SendInterval); err != nil {
		return nil, err
	}

	if err := encoder.Encode(n.ipv4); err != nil {
		return nil, err
	}

	if err := encoder.Encode(n.ipv6); err != nil {
		return nil, err
	}

	if err := encoder.Encode(n.Values); err != nil {
		return nil, err
	}

	return w.Bytes(), nil
}

// GobDecode gob interface
func (n *Node) GobDecode(buf []byte) error {
	n.mu.Lock()
	defer n.mu.Unlock()

	r := bytes.NewBuffer(buf)
	decoder := gob.NewDecoder(r)

	if err := decoder.Decode(&n.SendInterval); err != nil {
		return err
	}

	if err := decoder.Decode(&n.ipv4); err != nil {
		return err
	}

	if err := decoder.Decode(&n.ipv6); err != nil {
		return err
	}

	return decoder.Decode(&n.Values)
}

// Done returns a channel that can be used to wait till Multicast is stopped
func (n *Node) Done() <-chan struct{} {
	return n.mc.Done()
}

// RegisterCh returns a channel to know if the node should stay registered
func (n *Node) RegisterCh() <-chan struct{} {
	n.init()
	return n.registerCh
}

// KeepRegistered sends an anonymous struct{} to registeredChan to indicate the node should stay registered
func (n *Node) KeepRegistered() {
	n.init()
	n.registerCh <- struct{}{}
}

// Multicast start the mulicast ping
func (n *Node) Multicast(ctx context.Context, multicastAddress string) error {
	n.mu.Lock()
	n.ipv4 = localIPv4()
	n.ipv6 = localIPv6()

	if n.SendInterval.Seconds() == float64(0) {
		n.SendInterval = 1 * time.Second // default to 1 second
	}

	n.mu.Unlock()

	// Encode node to be sent via multicast
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	err := enc.Encode(n)
	if err != nil {
		return err
	}

	n.mu.Lock()
	defer n.mu.Unlock()
	n.mc = &multicast.Multicast{Address: multicastAddress}
	if err := n.mc.Send(ctx, n.SendInterval, buf.Bytes()); err != nil {
		return err
	}

	return nil
}

// Stop closes the StopCh to stop multicast sending
func (n *Node) Stop() {
	n.mu.Lock()
	defer n.mu.Unlock()
	n.mc.Stop()
}

// IPv4 getter for ipv4Address
func (n *Node) IPv4() net.IP {
	n.mu.Lock()
	defer n.mu.Unlock()
	return n.ipv4
}

// IPv6 getter for ipv6Address
func (n *Node) IPv6() net.IP {
	n.mu.Lock()
	defer n.mu.Unlock()
	return n.ipv6
}

// localIPv4 return the ipv4 address of the computer
// If it can't get the local ip it returns 127.0.0.1
// https://github.com/forestgiant/netutil
func localIPv4() net.IP {
	loopback := net.ParseIP("127.0.0.1")

	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return loopback
	}

	for _, addr := range addrs {
		// check the address type and make sure it's not loopback
		if ipnet, ok := addr.(*net.IPNet); ok {
			if !ipnet.IP.IsLoopback() {
				if ipnet.IP.To4() != nil {
					return ipnet.IP.To4()
				}
			}
		}
	}

	return loopback
}

// localIPv6 return the ipv6 address of the computer
// If it can't get the local ip it returns net.IPv6loopback
// https://github.com/forestgiant/netutil
func localIPv6() net.IP {
	loopback := net.IPv6loopback

	intfs, err := net.Interfaces()
	if err != nil {
		return loopback
	}

	for _, intf := range intfs {
		// If the interface is a loopback or doesn't have multicasting let's skip it
		if strings.Contains(intf.Flags.String(), net.FlagLoopback.String()) || !strings.Contains(intf.Flags.String(), net.FlagMulticast.String()) {
			continue
		}

		// Now let's check if the interface has an ipv6 address
		addrs, err := intf.Addrs()
		if err != nil {
			continue
		}

		for _, address := range addrs {
			if ipnet, ok := address.(*net.IPNet); ok {
				if !ipnet.IP.IsLoopback() {
					if ipnet.IP.To4() == nil {
						return ipnet.IP
					}
				}
			}
		}
	}

	return loopback
}
<|endoftext|>


<commit_before>package lib

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/namsral/flag"
)

// ReservedSpace -
const ReservedSpace = 450000000 // 450Mb

// Config -
type Config struct {
	DryRun         bool     `json:"dryRun"`
	NotifyPlan     int      `json:"notifyPlan"`
	NotifyTransfer int      `json:"notifyTransfer"`
	ReservedAmount int64    `json:"reservedAmount"`
	ReservedUnit   string   `json:"reservedUnit"`
	RsyncArgs      []string `json:"rsyncArgs"`
	Version        string   `json:"version"`
	Verbosity      int      `json:"verbosity"`
	CheckForUpdate int      `json:"checkForUpdate"`
}

// NotifyPlan/NotifyTransfer possible values
// 0 - no notification
// 1 - simple notification
// 2 - detailed notification

// Settings -
type Settings struct {
	Config

	Port       string
	LogDir     string
	APIFolders []string

	Location string
	confName string
}

const defaultConfLocation = "/boot/config/plugins/unbalance"

// NewSettings -
func NewSettings(name, version string, locations []string) (*Settings, error) {
	var port, logDir, folders, rsyncFlags, rsyncArgs, apiFolders string
	var dryRun bool
	var notifyCalc, notifyMove, notifyPlan, notifyTransfer, verbosity, checkForUpdate int

	flagset := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)

	flagset.StringVar(&port, "port", "6237", "port to run the server")
	flagset.StringVar(&logDir, "logdir", "/boot/logs", "pathname where log file will be written to")
	flagset.StringVar(&folders, "folders", "", "deprecated - do not use")
	flagset.BoolVar(&dryRun, "dryRun", true, "perform a dry-run rather than actual work")
	flagset.IntVar(&notifyCalc, "notifyCalc", 0, "deprecated - do not use") // deprecated
	flagset.IntVar(&notifyMove, "notifyMove", 0, "deprecated - do not use") // deprecated
	flagset.IntVar(&notifyPlan, "notifyPlan", 0, "notify via email after plan operation has completed (unraid notifications must be set up first): 0 - No notifications; 1 - Simple notifications; 2 - Detailed notifications")
	flagset.IntVar(&notifyTransfer, "notifyTransfer", 0, "notify via email after transfer operation has completed (unraid notifications must be set up first): 0 - No notifications; 1 - Simple notifications; 2 - Detailed notifications")
	flagset.StringVar(&rsyncFlags, "rsyncFlags", "", "deprecated - do not use") // deprecated
	flagset.StringVar(&rsyncArgs, "rsyncArgs", "", "custom rsync arguments")
	flagset.StringVar(&apiFolders, "apiFolders", "/var/local/emhttp", "folders to look for api endpoints")
	flagset.IntVar(&verbosity, "verbosity", 0, "include rsync output in log files: 0 (default) - include; 1 - do not include")
	flagset.IntVar(&checkForUpdate, "checkForUpdate", 1, "checkForUpdate: 0 - dont' check; 1 (default) - check")

	location := SearchFile(name, locations)
	if location != "" {
		flagset.String("config", filepath.Join(location, name), "config location")
	}

	if err := flagset.Parse(os.Args[1:]); err != nil {
		return nil, err
	}

	s := &Settings{}

	if rsyncArgs == "" {
		s.RsyncArgs = []string{""}
	} else {
		s.RsyncArgs = strings.Split(rsyncArgs, "|")
	}

	s.DryRun = dryRun
	s.NotifyPlan = notifyPlan
	s.NotifyTransfer = notifyTransfer
	s.ReservedAmount = ReservedSpace / 1000 / 1000
	s.ReservedUnit = "Mb"
	s.Verbosity = verbosity
	s.CheckForUpdate = checkForUpdate
	s.Version = version

	s.Port = port
	s.LogDir = logDir
	s.APIFolders = strings.Split(apiFolders, "|")
	s.Location = location
	s.confName = name

	return s, nil
}

// ToggleDryRun -
func (s *Settings) ToggleDryRun() {
	s.DryRun = !s.DryRun
}

// Save -
func (s *Settings) Save() (err error) {
	location := s.Location
	if location == "" {
		location = defaultConfLocation
	}

	confLocation := filepath.Join(location, s.confName)
	tmpFile := confLocation + ".tmp"

	if err = WriteLine(tmpFile, fmt.Sprintf("dryRun=%t", s.DryRun)); err != nil {
		return err
	}

	if err = WriteLine(tmpFile, fmt.Sprintf("notifyPlan=%d", s.NotifyPlan)); err != nil {
		return err
	}

	if err = WriteLine(tmpFile, fmt.Sprintf("notifyTransfer=%d", s.NotifyTransfer)); err != nil {
		return err
	}

	rsyncArgs := strings.Join(s.RsyncArgs, "|")
	if err = WriteLine(tmpFile, fmt.Sprintf("rsyncArgs=%s", rsyncArgs)); err != nil {
		return err
	}

	if err = WriteLine(tmpFile, fmt.Sprintf("verbosity=%d", s.Verbosity)); err != nil {
		return err
	}

	err = os.Rename(tmpFile, confLocation)
	if err != nil {
		return err
	}

	return
}
<commit_msg>Fix issue with empty rsync args settings<commit_after>package lib

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/namsral/flag"
)

// ReservedSpace -
const ReservedSpace = 450000000 // 450Mb

// Config -
type Config struct {
	DryRun         bool     `json:"dryRun"`
	NotifyPlan     int      `json:"notifyPlan"`
	NotifyTransfer int      `json:"notifyTransfer"`
	ReservedAmount int64    `json:"reservedAmount"`
	ReservedUnit   string   `json:"reservedUnit"`
	RsyncArgs      []string `json:"rsyncArgs"`
	Version        string   `json:"version"`
	Verbosity      int      `json:"verbosity"`
	CheckForUpdate int      `json:"checkForUpdate"`
}

// NotifyPlan/NotifyTransfer possible values
// 0 - no notification
// 1 - simple notification
// 2 - detailed notification

// Settings -
type Settings struct {
	Config

	Port       string
	LogDir     string
	APIFolders []string

	Location string
	confName string
}

const defaultConfLocation = "/boot/config/plugins/unbalance"

// NewSettings -
func NewSettings(name, version string, locations []string) (*Settings, error) {
	var port, logDir, folders, rsyncFlags, rsyncArgs, apiFolders string
	var dryRun bool
	var notifyCalc, notifyMove, notifyPlan, notifyTransfer, verbosity, checkForUpdate int

	flagset := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)

	flagset.StringVar(&port, "port", "6237", "port to run the server")
	flagset.StringVar(&logDir, "logdir", "/boot/logs", "pathname where log file will be written to")
	flagset.StringVar(&folders, "folders", "", "deprecated - do not use")
	flagset.BoolVar(&dryRun, "dryRun", true, "perform a dry-run rather than actual work")
	flagset.IntVar(&notifyCalc, "notifyCalc", 0, "deprecated - do not use") // deprecated
	flagset.IntVar(&notifyMove, "notifyMove", 0, "deprecated - do not use") // deprecated
	flagset.IntVar(&notifyPlan, "notifyPlan", 0, "notify via email after plan operation has completed (unraid notifications must be set up first): 0 - No notifications; 1 - Simple notifications; 2 - Detailed notifications")
	flagset.IntVar(&notifyTransfer, "notifyTransfer", 0, "notify via email after transfer operation has completed (unraid notifications must be set up first): 0 - No notifications; 1 - Simple notifications; 2 - Detailed notifications")
	flagset.StringVar(&rsyncFlags, "rsyncFlags", "", "deprecated - do not use") // deprecated
	flagset.StringVar(&rsyncArgs, "rsyncArgs", "", "custom rsync arguments")
	flagset.StringVar(&apiFolders, "apiFolders", "/var/local/emhttp", "folders to look for api endpoints")
	flagset.IntVar(&verbosity, "verbosity", 0, "include rsync output in log files: 0 (default) - include; 1 - do not include")
	flagset.IntVar(&checkForUpdate, "checkForUpdate", 1, "checkForUpdate: 0 - dont' check; 1 (default) - check")

	location := SearchFile(name, locations)
	if location != "" {
		flagset.String("config", filepath.Join(location, name), "config location")
	}

	if err := flagset.Parse(os.Args[1:]); err != nil {
		return nil, err
	}

	s := &Settings{}

	if rsyncArgs == "" {
		s.RsyncArgs = make([]string, 0)
	} else {
		s.RsyncArgs = strings.Split(rsyncArgs, "|")
	}

	s.DryRun = dryRun
	s.NotifyPlan = notifyPlan
	s.NotifyTransfer = notifyTransfer
	s.ReservedAmount = ReservedSpace / 1000 / 1000
	s.ReservedUnit = "Mb"
	s.Verbosity = verbosity
	s.CheckForUpdate = checkForUpdate
	s.Version = version

	s.Port = port
	s.LogDir = logDir
	s.APIFolders = strings.Split(apiFolders, "|")
	s.Location = location
	s.confName = name

	return s, nil
}

// ToggleDryRun -
func (s *Settings) ToggleDryRun() {
	s.DryRun = !s.DryRun
}

// Save -
func (s *Settings) Save() (err error) {
	location := s.Location
	if location == "" {
		location = defaultConfLocation
	}

	confLocation := filepath.Join(location, s.confName)
	tmpFile := confLocation + ".tmp"

	if err = WriteLine(tmpFile, fmt.Sprintf("dryRun=%t", s.DryRun)); err != nil {
		return err
	}

	if err = WriteLine(tmpFile, fmt.Sprintf("notifyPlan=%d", s.NotifyPlan)); err != nil {
		return err
	}

	if err = WriteLine(tmpFile, fmt.Sprintf("notifyTransfer=%d", s.NotifyTransfer)); err != nil {
		return err
	}

	rsyncArgs := strings.Join(s.RsyncArgs, "|")
	if err = WriteLine(tmpFile, fmt.Sprintf("rsyncArgs=%s", rsyncArgs)); err != nil {
		return err
	}

	if err = WriteLine(tmpFile, fmt.Sprintf("verbosity=%d", s.Verbosity)); err != nil {
		return err
	}

	err = os.Rename(tmpFile, confLocation)
	if err != nil {
		return err
	}

	return
}
<|endoftext|>