<commit_before>// Copyright 2017 Vector Creations Ltd
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package routing

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/matrix-org/dendrite/clientapi/httputil"
	"github.com/matrix-org/dendrite/clientapi/jsonerror"
	"github.com/matrix-org/dendrite/clientapi/producers"
	"github.com/matrix-org/dendrite/common/config"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
)

// Send implements /_matrix/federation/v1/send/{txnID}
func Send(
	httpReq *http.Request,
	request *gomatrixserverlib.FederationRequest,
	txnID gomatrixserverlib.TransactionID,
	cfg config.Dendrite,
	query api.RoomserverQueryAPI,
	producer *producers.RoomserverProducer,
	keys gomatrixserverlib.KeyRing,
	federation *gomatrixserverlib.FederationClient,
) util.JSONResponse {

	t := txnReq{
		context:    httpReq.Context(),
		query:      query,
		producer:   producer,
		keys:       keys,
		federation: federation,
	}
	if err := json.Unmarshal(request.Content(), &t); err != nil {
		return util.JSONResponse{
			Code: 400,
			JSON: jsonerror.NotJSON("The request body could not be decoded into valid JSON. " + err.Error()),
		}
	}

	t.Origin = request.Origin()
	t.TransactionID = txnID
	t.Destination = cfg.Matrix.ServerName

	resp, err := t.processTransaction()
	if err != nil {
		return httputil.LogThenError(httpReq, err)
	}

	return util.JSONResponse{
		Code: 200,
		JSON: resp,
	}
}

type txnReq struct {
	gomatrixserverlib.Transaction
	context    context.Context
	query      api.RoomserverQueryAPI
	producer   *producers.RoomserverProducer
	keys       gomatrixserverlib.KeyRing
	federation *gomatrixserverlib.FederationClient
}

func (t *txnReq) processTransaction() (*gomatrixserverlib.RespSend, error) {
	// Check the event signatures
	if err := gomatrixserverlib.VerifyEventSignatures(t.context, t.PDUs, t.keys); err != nil {
		return nil, err
	}

	// Process the events.
	results := map[string]gomatrixserverlib.PDUResult{}
	for _, e := range t.PDUs {
		err := t.processEvent(e)
		if err != nil {
			// If the error is due to the event itself being bad then we skip
			// it and move onto the next event. We report an error so that the
			// sender knows that we have skipped processing it.
			//
			// However if the event is due to a temporary failure in our server
			// such as a database being unavailable then we should bail, and
			// hope that the sender will retry when we are feeling better.
			//
			// It is uncertain what we should do if an event fails because
			// we failed to fetch more information from the sending server.
			// For example if a request to /state fails.
			// If we skip the event then we risk missing the event until we
			// receive another event referencing it.
			// If we bail and stop processing then we risk wedging incoming
			// transactions from that server forever.
			switch err.(type) {
			case unknownRoomError:
			case *gomatrixserverlib.NotAllowed:
			default:
				// Any other error should be the result of a temporary error in
				// our server so we should bail processing the transaction entirely.
				return nil, err
			}
			results[e.EventID()] = gomatrixserverlib.PDUResult{
				Error: err.Error(),
			}
		} else {
			results[e.EventID()] = gomatrixserverlib.PDUResult{}
		}
	}

	// TODO: Process the EDUs.

	return &gomatrixserverlib.RespSend{PDUs: results}, nil
}

type unknownRoomError struct {
	roomID string
}

func (e unknownRoomError) Error() string { return fmt.Sprintf("unknown room %q", e.roomID) }

func (t *txnReq) processEvent(e gomatrixserverlib.Event) error {
	prevEventIDs := e.PrevEventIDs()

	// Fetch the state needed to authenticate the event.
	needed := gomatrixserverlib.StateNeededForAuth([]gomatrixserverlib.Event{e})
	stateReq := api.QueryStateAfterEventsRequest{
		RoomID:       e.RoomID(),
		PrevEventIDs: prevEventIDs,
		StateToFetch: needed.Tuples(),
	}
	var stateResp api.QueryStateAfterEventsResponse
	if err := t.query.QueryStateAfterEvents(t.context, &stateReq, &stateResp); err != nil {
		return err
	}

	if !stateResp.RoomExists {
		// TODO: When synapse receives a message for a room it is not in it
		// asks the remote server for the state of the room so that it can
		// check if the remote server knows of a join "m.room.member" event
		// that this server is unaware of.
		// However generally speaking we should reject events for rooms we
		// aren't a member of.
		return unknownRoomError{e.RoomID()}
	}

	if !stateResp.PrevEventsExist {
		return t.processEventWithMissingState(e)
	}

	// Check that the event is allowed by the state at the event.
	if err := checkAllowedByState(e, stateResp.StateEvents); err != nil {
		return err
	}

	// TODO: Check that the roomserver has a copy of all of the auth_events.
	// TODO: Check that the event is allowed by its auth_events.

	// pass the event to the roomserver
	return t.producer.SendEvents(t.context, []gomatrixserverlib.Event{e}, api.DoNotSendToOtherServers, nil)
}

func checkAllowedByState(e gomatrixserverlib.Event, stateEvents []gomatrixserverlib.Event) error {
	authUsingState := gomatrixserverlib.NewAuthEvents(nil)
	for i := range stateEvents {
		err := authUsingState.AddEvent(&stateEvents[i])
		if err != nil {
			return err
		}
	}
	return gomatrixserverlib.Allowed(e, &authUsingState)
}

func (t *txnReq) processEventWithMissingState(e gomatrixserverlib.Event) error {
	// We are missing the previous events for this events.
	// This means that there is a gap in our view of the history of the
	// room. There two ways that we can handle such a gap:
	//   1) We can fill in the gap using /get_missing_events
	//   2) We can leave the gap and request the state of the room at
	//      this event from the remote server using either /state_ids
	//      or /state.
	// Synapse will attempt to do 1 and if that fails or if the gap is
	// too large then it will attempt 2.
	// Synapse will use /state_ids if possible since usually the state
	// is largely unchanged and it is more efficient to fetch a list of
	// event ids and then use /event to fetch the individual events.
	// However not all version of synapse support /state_ids so you may
	// need to fallback to /state.
	// TODO: Attempt to fill in the gap using /get_missing_events
	// TODO: Attempt to fetch the state using /state_ids and /events
	state, err := t.federation.LookupState(t.context, t.Origin, e.RoomID(), e.EventID())
	if err != nil {
		return err
	}
	// Check that the returned state is valid.
	if err := state.Check(t.context, t.keys); err != nil {
		return err
	}
	// Check that the event is allowed by the state.
	if err := checkAllowedByState(e, state.StateEvents); err != nil {
		return err
	}
	// pass the event along with the state to the roomserver
	return t.producer.SendEventWithState(t.context, state, e)
}
<commit_msg>s/VerifyEventSignatures/VerifyAllEventSignatures/<commit_after>// Copyright 2017 Vector Creations Ltd
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package routing

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/matrix-org/dendrite/clientapi/httputil"
	"github.com/matrix-org/dendrite/clientapi/jsonerror"
	"github.com/matrix-org/dendrite/clientapi/producers"
	"github.com/matrix-org/dendrite/common/config"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
)

// Send implements /_matrix/federation/v1/send/{txnID}
func Send(
	httpReq *http.Request,
	request *gomatrixserverlib.FederationRequest,
	txnID gomatrixserverlib.TransactionID,
	cfg config.Dendrite,
	query api.RoomserverQueryAPI,
	producer *producers.RoomserverProducer,
	keys gomatrixserverlib.KeyRing,
	federation *gomatrixserverlib.FederationClient,
) util.JSONResponse {

	t := txnReq{
		context:    httpReq.Context(),
		query:      query,
		producer:   producer,
		keys:       keys,
		federation: federation,
	}
	if err := json.Unmarshal(request.Content(), &t); err != nil {
		return util.JSONResponse{
			Code: 400,
			JSON: jsonerror.NotJSON("The request body could not be decoded into valid JSON. " + err.Error()),
		}
	}

	t.Origin = request.Origin()
	t.TransactionID = txnID
	t.Destination = cfg.Matrix.ServerName

	resp, err := t.processTransaction()
	if err != nil {
		return httputil.LogThenError(httpReq, err)
	}

	return util.JSONResponse{
		Code: 200,
		JSON: resp,
	}
}

type txnReq struct {
	gomatrixserverlib.Transaction
	context    context.Context
	query      api.RoomserverQueryAPI
	producer   *producers.RoomserverProducer
	keys       gomatrixserverlib.KeyRing
	federation *gomatrixserverlib.FederationClient
}

func (t *txnReq) processTransaction() (*gomatrixserverlib.RespSend, error) {
	// Check the event signatures
	if err := gomatrixserverlib.VerifyAllEventSignatures(t.context, t.PDUs, t.keys); err != nil {
		return nil, err
	}

	// Process the events.
	results := map[string]gomatrixserverlib.PDUResult{}
	for _, e := range t.PDUs {
		err := t.processEvent(e)
		if err != nil {
			// If the error is due to the event itself being bad then we skip
			// it and move onto the next event. We report an error so that the
			// sender knows that we have skipped processing it.
			//
			// However if the event is due to a temporary failure in our server
			// such as a database being unavailable then we should bail, and
			// hope that the sender will retry when we are feeling better.
			//
			// It is uncertain what we should do if an event fails because
			// we failed to fetch more information from the sending server.
			// For example if a request to /state fails.
			// If we skip the event then we risk missing the event until we
			// receive another event referencing it.
			// If we bail and stop processing then we risk wedging incoming
			// transactions from that server forever.
			switch err.(type) {
			case unknownRoomError:
			case *gomatrixserverlib.NotAllowed:
			default:
				// Any other error should be the result of a temporary error in
				// our server so we should bail processing the transaction entirely.
				return nil, err
			}
			results[e.EventID()] = gomatrixserverlib.PDUResult{
				Error: err.Error(),
			}
		} else {
			results[e.EventID()] = gomatrixserverlib.PDUResult{}
		}
	}

	// TODO: Process the EDUs.

	return &gomatrixserverlib.RespSend{PDUs: results}, nil
}

type unknownRoomError struct {
	roomID string
}

func (e unknownRoomError) Error() string { return fmt.Sprintf("unknown room %q", e.roomID) }

func (t *txnReq) processEvent(e gomatrixserverlib.Event) error {
	prevEventIDs := e.PrevEventIDs()

	// Fetch the state needed to authenticate the event.
	needed := gomatrixserverlib.StateNeededForAuth([]gomatrixserverlib.Event{e})
	stateReq := api.QueryStateAfterEventsRequest{
		RoomID:       e.RoomID(),
		PrevEventIDs: prevEventIDs,
		StateToFetch: needed.Tuples(),
	}
	var stateResp api.QueryStateAfterEventsResponse
	if err := t.query.QueryStateAfterEvents(t.context, &stateReq, &stateResp); err != nil {
		return err
	}

	if !stateResp.RoomExists {
		// TODO: When synapse receives a message for a room it is not in it
		// asks the remote server for the state of the room so that it can
		// check if the remote server knows of a join "m.room.member" event
		// that this server is unaware of.
		// However generally speaking we should reject events for rooms we
		// aren't a member of.
		return unknownRoomError{e.RoomID()}
	}

	if !stateResp.PrevEventsExist {
		return t.processEventWithMissingState(e)
	}

	// Check that the event is allowed by the state at the event.
	if err := checkAllowedByState(e, stateResp.StateEvents); err != nil {
		return err
	}

	// TODO: Check that the roomserver has a copy of all of the auth_events.
	// TODO: Check that the event is allowed by its auth_events.

	// pass the event to the roomserver
	return t.producer.SendEvents(t.context, []gomatrixserverlib.Event{e}, api.DoNotSendToOtherServers, nil)
}

func checkAllowedByState(e gomatrixserverlib.Event, stateEvents []gomatrixserverlib.Event) error {
	authUsingState := gomatrixserverlib.NewAuthEvents(nil)
	for i := range stateEvents {
		err := authUsingState.AddEvent(&stateEvents[i])
		if err != nil {
			return err
		}
	}
	return gomatrixserverlib.Allowed(e, &authUsingState)
}

func (t *txnReq) processEventWithMissingState(e gomatrixserverlib.Event) error {
	// We are missing the previous events for this events.
	// This means that there is a gap in our view of the history of the
	// room. There two ways that we can handle such a gap:
	//   1) We can fill in the gap using /get_missing_events
	//   2) We can leave the gap and request the state of the room at
	//      this event from the remote server using either /state_ids
	//      or /state.
	// Synapse will attempt to do 1 and if that fails or if the gap is
	// too large then it will attempt 2.
	// Synapse will use /state_ids if possible since usually the state
	// is largely unchanged and it is more efficient to fetch a list of
	// event ids and then use /event to fetch the individual events.
	// However not all version of synapse support /state_ids so you may
	// need to fallback to /state.
	// TODO: Attempt to fill in the gap using /get_missing_events
	// TODO: Attempt to fetch the state using /state_ids and /events
	state, err := t.federation.LookupState(t.context, t.Origin, e.RoomID(), e.EventID())
	if err != nil {
		return err
	}
	// Check that the returned state is valid.
	if err := state.Check(t.context, t.keys); err != nil {
		return err
	}
	// Check that the event is allowed by the state.
	if err := checkAllowedByState(e, state.StateEvents); err != nil {
		return err
	}
	// pass the event along with the state to the roomserver
	return t.producer.SendEventWithState(t.context, state, e)
}
<|endoftext|>


<commit_before>// Package wikipedia implements a Service which adds !commands for Wikipedia search.
package wikipedia

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

	log "github.com/Sirupsen/logrus"
	"github.com/jaytaylor/html2text"
	"github.com/matrix-org/go-neb/types"
	"github.com/matrix-org/gomatrix"
)

// ServiceType of the Wikipedia service
const ServiceType = "wikipedia"
const maxExtractLength = 1024 // Max length of extract string in bytes

var httpClient = &http.Client{}

// Search results (returned by search query)
type wikipediaSearchResults struct {
	Query wikipediaQuery `json:"query"`
}

// Wikipeda pages returned in search results
type wikipediaQuery struct {
	Pages map[string]wikipediaPage `json:"pages"`
}

// Representation of an individual wikipedia page
type wikipediaPage struct {
	PageID    int64  `json:"pageid"`
	NS        int    `json:"ns"`
	Title     string `json:"title"`
	Touched   string `json:"touched"`
	LastRevID int64  `json:"lastrevid"`
	Extract   string `json:"extract"`
}

// Service contains the Config fields for the Wikipedia service.
type Service struct {
	types.DefaultService
}

// Commands supported:
//    !wikipedia some_search_query_without_quotes
// Responds with a suitable article extract and link to the referenced page into the same room as the command.
func (s *Service) Commands(client *gomatrix.Client) []types.Command {
	return []types.Command{
		types.Command{
			Path: []string{"wikipedia"},
			Command: func(roomID, userID string, args []string) (interface{}, error) {
				return s.cmdWikipediaSearch(client, roomID, userID, args)
			},
		},
	}
}

// usageMessage returns a matrix TextMessage representation of the service usage
func usageMessage() *gomatrix.TextMessage {
	return &gomatrix.TextMessage{"m.notice",
		`Usage: !wikipedia search_text`}
}

func (s *Service) cmdWikipediaSearch(client *gomatrix.Client, roomID, userID string, args []string) (interface{}, error) {
	// Check for query text
	if len(args) < 1 {
		return usageMessage(), nil
	}

	// Get the query text and per,form search
	querySentence := strings.Join(args, " ")
	searchResultPage, err := s.text2Wikipedia(querySentence)
	if err != nil {
		return nil, err
	}

	// No article extracts
	if searchResultPage == nil || searchResultPage.Extract == "" {
		return gomatrix.TextMessage{
			MsgType: "m.notice",
			Body:    "No results",
		}, nil
	}

	// Convert article HTML to text
	extractText, err := html2text.FromString(searchResultPage.Extract)
	if err != nil {
		return gomatrix.TextMessage{
			MsgType: "m.notice",
			Body:    "Failed to convert extract to plain text - " + err.Error(),
		}, nil
	}

	// Truncate the extract text, if necessary
	if len(extractText) > maxExtractLength {
		extractText = extractText[:maxExtractLength] + "..."
	}

	// Add a link to the bottom of the extract
	extractText += fmt.Sprintf("\nhttp://en.wikipedia.org/?curid=%d", searchResultPage.PageID)

	// Return article extract
	return gomatrix.TextMessage{
		MsgType: "m.notice",
		Body:    extractText,
	}, nil
}

// text2Wikipedia returns a Wikipedia article summary
func (s *Service) text2Wikipedia(query string) (*wikipediaPage, error) {
	log.Info("Searching Wikipedia for: ", query)

	u, err := url.Parse("https://en.wikipedia.org/w/api.php")
	if err != nil {
		return nil, err
	}

	// Example query - https://en.wikipedia.org/w/api.php?action=query&prop=extracts&format=json&exintro=&titles=RMS+Titanic
	q := u.Query()
	q.Set("action", "query")  // Action - query for articles
	q.Set("prop", "extracts") // Return article extracts
	q.Set("format", "json")
	q.Set("redirects", "")
	// q.Set("exintro", "")
	q.Set("titles", query) // Text to search for

	u.RawQuery = q.Encode()
	// log.Info("Request URL: ", u)

	// Perform wikipedia search request
	res, err := httpClient.Get(u.String())
	if res != nil {
		defer res.Body.Close()
	}
	if err != nil {
		return nil, err
	}
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		return nil, fmt.Errorf("Request error: %d, %s", res.StatusCode, response2String(res))
	}

	// Parse search results
	var searchResults wikipediaSearchResults
	// log.Info(response2String(res))
	if err := json.NewDecoder(res.Body).Decode(&searchResults); err != nil {
		return nil, fmt.Errorf("ERROR - %s", err.Error())
	} else if len(searchResults.Query.Pages) < 1 {
		return nil, fmt.Errorf("No articles found")
	}

	// Return only the first search result with an extract
	for _, page := range searchResults.Query.Pages {
		if page.Extract != "" {
			return &page, nil
		}
	}

	return nil, fmt.Errorf("No articles with extracts found")
}

// response2String returns a string representation of an HTTP response body
func response2String(res *http.Response) string {
	bs, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "Failed to decode response body"
	}
	str := string(bs)
	return str
}

// Initialise the service
func init() {
	types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service {
		return &Service{
			DefaultService: types.NewDefaultService(serviceID, serviceUserID, ServiceType),
		}
	})
}
<commit_msg>Comment struct fields (specifically map keys)<commit_after>// Package wikipedia implements a Service which adds !commands for Wikipedia search.
package wikipedia

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

	log "github.com/Sirupsen/logrus"
	"github.com/jaytaylor/html2text"
	"github.com/matrix-org/go-neb/types"
	"github.com/matrix-org/gomatrix"
)

// ServiceType of the Wikipedia service
const ServiceType = "wikipedia"
const maxExtractLength = 1024 // Max length of extract string in bytes

var httpClient = &http.Client{}

// Search results (returned by search query)
type wikipediaSearchResults struct {
	Query wikipediaQuery `json:"query"` // Containter for the query response
}

// Wikipeda pages returned in search results
type wikipediaQuery struct {
	Pages map[string]wikipediaPage `json:"pages"` // Map of wikipedia page IDs to page objects
}

// Representation of an individual wikipedia page
type wikipediaPage struct {
	PageID    int64  `json:"pageid"`    // Unique ID for the wikipedia page
	NS        int    `json:"ns"`        // Namespace ID
	Title     string `json:"title"`     // Page title text
	Touched   string `json:"touched"`   // Date that the page was last touched / modified
	LastRevID int64  `json:"lastrevid"` //
	Extract   string `json:"extract"`   // Page extract text
}

// Service contains the Config fields for the Wikipedia service.
type Service struct {
	types.DefaultService
}

// Commands supported:
//    !wikipedia some_search_query_without_quotes
// Responds with a suitable article extract and link to the referenced page into the same room as the command.
func (s *Service) Commands(client *gomatrix.Client) []types.Command {
	return []types.Command{
		types.Command{
			Path: []string{"wikipedia"},
			Command: func(roomID, userID string, args []string) (interface{}, error) {
				return s.cmdWikipediaSearch(client, roomID, userID, args)
			},
		},
	}
}

// usageMessage returns a matrix TextMessage representation of the service usage
func usageMessage() *gomatrix.TextMessage {
	return &gomatrix.TextMessage{"m.notice",
		`Usage: !wikipedia search_text`}
}

func (s *Service) cmdWikipediaSearch(client *gomatrix.Client, roomID, userID string, args []string) (interface{}, error) {
	// Check for query text
	if len(args) < 1 {
		return usageMessage(), nil
	}

	// Get the query text and per,form search
	querySentence := strings.Join(args, " ")
	searchResultPage, err := s.text2Wikipedia(querySentence)
	if err != nil {
		return nil, err
	}

	// No article extracts
	if searchResultPage == nil || searchResultPage.Extract == "" {
		return gomatrix.TextMessage{
			MsgType: "m.notice",
			Body:    "No results",
		}, nil
	}

	// Convert article HTML to text
	extractText, err := html2text.FromString(searchResultPage.Extract)
	if err != nil {
		return gomatrix.TextMessage{
			MsgType: "m.notice",
			Body:    "Failed to convert extract to plain text - " + err.Error(),
		}, nil
	}

	// Truncate the extract text, if necessary
	if len(extractText) > maxExtractLength {
		extractText = extractText[:maxExtractLength] + "..."
	}

	// Add a link to the bottom of the extract
	extractText += fmt.Sprintf("\nhttp://en.wikipedia.org/?curid=%d", searchResultPage.PageID)

	// Return article extract
	return gomatrix.TextMessage{
		MsgType: "m.notice",
		Body:    extractText,
	}, nil
}

// text2Wikipedia returns a Wikipedia article summary
func (s *Service) text2Wikipedia(query string) (*wikipediaPage, error) {
	log.Info("Searching Wikipedia for: ", query)

	u, err := url.Parse("https://en.wikipedia.org/w/api.php")
	if err != nil {
		return nil, err
	}

	// Example query - https://en.wikipedia.org/w/api.php?action=query&prop=extracts&format=json&exintro=&titles=RMS+Titanic
	q := u.Query()
	q.Set("action", "query")  // Action - query for articles
	q.Set("prop", "extracts") // Return article extracts
	q.Set("format", "json")
	q.Set("redirects", "")
	// q.Set("exintro", "")
	q.Set("titles", query) // Text to search for

	u.RawQuery = q.Encode()
	// log.Info("Request URL: ", u)

	// Perform wikipedia search request
	res, err := httpClient.Get(u.String())
	if res != nil {
		defer res.Body.Close()
	}
	if err != nil {
		return nil, err
	}
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		return nil, fmt.Errorf("Request error: %d, %s", res.StatusCode, response2String(res))
	}

	// Parse search results
	var searchResults wikipediaSearchResults
	// log.Info(response2String(res))
	if err := json.NewDecoder(res.Body).Decode(&searchResults); err != nil {
		return nil, fmt.Errorf("ERROR - %s", err.Error())
	} else if len(searchResults.Query.Pages) < 1 {
		return nil, fmt.Errorf("No articles found")
	}

	// Return only the first search result with an extract
	for _, page := range searchResults.Query.Pages {
		if page.Extract != "" {
			return &page, nil
		}
	}

	return nil, fmt.Errorf("No articles with extracts found")
}

// response2String returns a string representation of an HTTP response body
func response2String(res *http.Response) string {
	bs, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return "Failed to decode response body"
	}
	str := string(bs)
	return str
}

// Initialise the service
func init() {
	types.RegisterService(func(serviceID, serviceUserID, webhookEndpointURL string) types.Service {
		return &Service{
			DefaultService: types.NewDefaultService(serviceID, serviceUserID, ServiceType),
		}
	})
}
<|endoftext|>


<commit_before>package db

import (
	"github.com/stellar/go-stellar-base/xdr"
	"golang.org/x/net/context"
)

const OrderBookSummaryPageSize = 20

const OrderBookSummarySQL = `
SELECT
  *,
	(pricen :: double precision / priced :: double precision) as pricef

FROM
((
	SELECT
		'ask' as type,
		co.pricen,
		co.priced,
		SUM(co.amount) as amount

	FROM  offers co

	WHERE co.sellingassettype = $1
	AND   co.sellingassetcode = $2
	AND   co.sellingissuer    = $3
	AND   co.buyingassettype  = $4
	AND   co.buyingassetcode  = $5
	AND   co.buyingissuer     = $6

	GROUP BY
		co.pricen,
		co.priced,
		co.price
	LIMIT 15

) UNION (
	SELECT
		'bid'  as type,
		co.priced as pricen,
		co.pricen as priced,
		SUM(co.amount) as amount

	FROM offers co

	WHERE co.sellingassettype = $4
	AND   co.sellingassetcode = $5
	AND   co.sellingissuer    = $6
	AND   co.buyingassettype  = $1
	AND   co.buyingassetcode  = $2
	AND   co.buyingissuer     = $3

	GROUP BY
		co.pricen,
		co.priced,
		co.price
	LIMIT 15
)) summary

ORDER BY type, pricef
`

type OrderBookSummaryQuery struct {
	SqlQuery
	BaseType      xdr.AssetType
	BaseCode      string
	BaseIssuer    string
	CounterType   xdr.AssetType
	CounterCode   string
	CounterIssuer string
}

func (q OrderBookSummaryQuery) Invert() OrderBookSummaryQuery {

	return OrderBookSummaryQuery{
		SqlQuery:      q.SqlQuery,
		BaseType:      q.CounterType,
		BaseCode:      q.CounterCode,
		BaseIssuer:    q.CounterIssuer,
		CounterType:   q.BaseType,
		CounterCode:   q.BaseCode,
		CounterIssuer: q.BaseIssuer,
	}
}

func (q OrderBookSummaryQuery) Select(ctx context.Context, dest interface{}) error {
	args := []interface{}{
		q.BaseType,
		q.BaseCode,
		q.BaseIssuer,
		q.CounterType,
		q.CounterCode,
		q.CounterIssuer,
	}

	return q.SqlQuery.SelectRaw(ctx, OrderBookSummarySQL, args, dest)
}
<commit_msg>Add sql documentation for order book summary query<commit_after>package db

import (
	"github.com/stellar/go-stellar-base/xdr"
	"golang.org/x/net/context"
)

const OrderBookSummaryPageSize = 20

const OrderBookSummarySQL = `
SELECT
	*,
	(pricen :: double precision / priced :: double precision) as pricef

FROM
((
	-- This query returns the "asks" portion of the summary, and it is very straightforward
	SELECT
		'ask' as type,
		co.pricen,
		co.priced,
		SUM(co.amount) as amount

	FROM  offers co

	WHERE co.sellingassettype = $1
	AND   co.sellingassetcode = $2
	AND   co.sellingissuer    = $3
	AND   co.buyingassettype  = $4
	AND   co.buyingassetcode  = $5
	AND   co.buyingissuer     = $6

	GROUP BY
		co.pricen,
		co.priced,
		co.price
	LIMIT 15

) UNION (
	-- This query returns the "bids" portion, inverting the where clauses
	-- and the pricen/priced.  This inversion is necessary to produce the "bid"
	-- view of a given offer (which are stored in the db as an offer to sell)
	SELECT
		'bid'  as type,
		co.priced as pricen,
		co.pricen as priced,
		SUM(co.amount) as amount

	FROM offers co

	WHERE co.sellingassettype = $4
	AND   co.sellingassetcode = $5
	AND   co.sellingissuer    = $6
	AND   co.buyingassettype  = $1
	AND   co.buyingassetcode  = $2
	AND   co.buyingissuer     = $3

	GROUP BY
		co.pricen,
		co.priced,
		co.price
	LIMIT 15
)) summary

ORDER BY type, pricef
`

type OrderBookSummaryQuery struct {
	SqlQuery
	BaseType      xdr.AssetType
	BaseCode      string
	BaseIssuer    string
	CounterType   xdr.AssetType
	CounterCode   string
	CounterIssuer string
}

func (q OrderBookSummaryQuery) Invert() OrderBookSummaryQuery {

	return OrderBookSummaryQuery{
		SqlQuery:      q.SqlQuery,
		BaseType:      q.CounterType,
		BaseCode:      q.CounterCode,
		BaseIssuer:    q.CounterIssuer,
		CounterType:   q.BaseType,
		CounterCode:   q.BaseCode,
		CounterIssuer: q.BaseIssuer,
	}
}

func (q OrderBookSummaryQuery) Select(ctx context.Context, dest interface{}) error {
	args := []interface{}{
		q.BaseType,
		q.BaseCode,
		q.BaseIssuer,
		q.CounterType,
		q.CounterCode,
		q.CounterIssuer,
	}

	return q.SqlQuery.SelectRaw(ctx, OrderBookSummarySQL, args, dest)
}
<|endoftext|>


<commit_before>package core

import (
	"database/sql"
	"errors"

	_ "github.com/go-sql-driver/mysql"

	"../utils"
)

/*
	Container info
*/
type Container struct {
	ID         string
	ProbeID    string
	Hostname   string
	Image      string
	IPAddress  string
	MacAddress string
}

/*
	Container's stats
*/
type Stat struct {
	ContainerID string
	Time        int64
	SizeRootFs  uint64
	SizeRw      uint64
	SizeMemory  uint64
	Running     bool
}

/*
	SQL Variables
*/
var (
	DB *sql.DB // DB

	// Prepared queries
	GetProbeIDStmt            *sql.Stmt
	InsertProbeStmt           *sql.Stmt
	InsertContainerStmt       *sql.Stmt
	DeleteContainerStmt       *sql.Stmt
	InsertStatStmt            *sql.Stmt
	DeleteStatStmt            *sql.Stmt
	GetContainerByIdStmt      *sql.Stmt
	GetStatsByContainerIdStmt *sql.Stmt
)

func InitSQL() {
	var err error // Error handling

	// Connect DB
	sqlc := DGConfig.DockerGuard.SQLServer // SQL Config
	DB, err = sql.Open("mysql", sqlc.User+":"+sqlc.Pass+"@tcp("+sqlc.IP+":"+utils.I2S(sqlc.Port)+")/"+sqlc.DB)
	if err != nil {
		l.Critical(err)
	}
	err = DB.Ping()
	if err != nil {
		l.Critical(err)
	}
	l.Verbose("Connected to SQL database")

	// Prepare DB queries
	GetProbeIDStmt, err = DB.Prepare("SELECT id FROM probes WHERE name=?")
	if err != nil {
		l.Critical("Can't create InsertContainerStmt:", err)
	}
	InsertProbeStmt, err = DB.Prepare("INSERT INTO probes VALUES (DEFAULT,?)")
	if err != nil {
		l.Critical("Can't create InsertContainerStmt:", err)
	}
	InsertContainerStmt, err = DB.Prepare("INSERT INTO containers VALUES (?,?,?,?,?,?)")
	if err != nil {
		l.Critical("Can't create InsertContainerStmt:", err)
	}
	DeleteContainerStmt, err = DB.Prepare("DELETE FROM containers WHERE id=?")
	if err != nil {
		l.Critical("Can't create DeleteContainerStmt:", err)
	}
	InsertStatStmt, err = DB.Prepare("INSERT INTO stats VALUES (?,?,?,?,?,?)")
	if err != nil {
		l.Critical("Can't create InsertStatStmt:", err)
	}
	DeleteStatStmt, err = DB.Prepare("DELETE FROM stats WHERE containerid=? AND time=?")
	if err != nil {
		l.Critical("Can't create DeleteStatStmt:", err)
	}
	GetContainerByIdStmt, err = DB.Prepare("SELECT * FROM containers WHERE id=?")
	if err != nil {
		l.Critical("Can't create GetContainerByIdStmt:", err)
	}
	GetStatsByContainerIdStmt, err = DB.Prepare("SELECT * FROM stats WHERE containerid=?")
	if err != nil {
		l.Critical("Can't create DeleteStatStmt:", err)
	}
}

/*
	Get a probe ID from sql server
	If does not exist: create it and return the probe ID
*/
func GetProbeID(name string) (int, error) {
	var id int64  // Probe ID to return
	var err error // Error handling

	err = GetProbeIDStmt.QueryRow(name).Scan(&id)
	if err != nil {
		if err.Error() == "sql: no rows in result set" {
			result, err := InsertProbeStmt.Exec(name)
			if err != nil {
				l.Error("GetProbeID:", err)
				return 0, err
			}
			id, err = result.LastInsertId()
			if err != nil {
				l.Error("GetProbeID:", err)
				return 0, err
			}
		} else {
			l.Error("GetProbeID:", err)
			return 0, err
		}
	}

	return int(id), nil
}

/*
	Insert a container
*/
func (c *Container) Insert() error {
	var err error // Error handling

	_, err = InsertContainerStmt.Exec(c.ID, c.ProbeID, c.Hostname, c.Image, c.IPAddress, c.MacAddress)

	return err
}

/*
	Delete a container
*/
func (c *Container) Delete() error {
	var err error // Error handling

	_, err = DeleteContainerStmt.Exec(c.ID)

	return err
}

/*
	Get containers by $field
*/
func GetContainersBy(field string, value interface{}) ([]Container, error) {
	var containers []Container // Containers to return
	var rows *sql.Rows         // SQL Rows
	var err error              // Error handling

	// Protection against SQL injection
	var fieldExists bool = false
	for _, i := range []string{"id", "hostname", "image", "ip", "mac"} {
		if field == i {
			fieldExists = true
		}
	}
	if !fieldExists {
		l.Error("GetContainersBy: Field (" + field + ") does not exist.")
		return containers, errors.New("GetContainersBy: Field (" + field + ") does not exist.")
	}

	// Get containers
	rows, err = DB.Query("SELECT * FROM containers WHERE "+field+"=?", value)
	if err != nil {
		l.Error("GetContainersBy:", err)
		return containers, err
	}
	defer rows.Close()
	for rows.Next() {
		var tmpContainer Container

		if err = rows.Scan(&tmpContainer.ID,
			&tmpContainer.ProbeID,
			&tmpContainer.Hostname,
			&tmpContainer.Image,
			&tmpContainer.IPAddress,
			&tmpContainer.MacAddress); err != nil {
			l.Error("GetContainersBy:", err)
			return containers, err
		}

		containers = append(containers, tmpContainer)
	}
	if err = rows.Err(); err != nil {
		l.Error("GetContainersBy:", err)
		return containers, err
	}

	return containers, nil
}

/*
	Get a container by id
*/
func GetContainerById(id string) (Container, error) {
	var container Container // Container to return
	var err error           // Error handling

	err = GetContainerByIdStmt.QueryRow(id).Scan(&container.ID, &container.ProbeID, &container.Hostname, &container.Image, &container.IPAddress, &container.MacAddress)
	if err != nil {
		l.Error("GetContainerById:", err)
		return container, err
	}

	return container, nil
}

/*
	Insert a stat
*/
func (s *Stat) Insert() error {
	var err error // Error handling

	_, err = InsertStatStmt.Exec(s.ContainerID, s.Time, s.SizeRootFs, s.SizeRw, s.SizeMemory, s.Running)

	return err
}

/*
	Delete a stat
*/
func (s *Stat) Delete() error {
	var err error // Error handling

	_, err = DeleteStatStmt.Exec(s.ContainerID, s.Time)

	return err
}

/*
	Get stats by container id
*/
func GetStatsByContainerID(containerID string) ([]Stat, error) {
	var stats []Stat   // List of stats to return
	var err error      // Erro handling
	var tmpStat Stat   // Temporary stat
	var rows *sql.Rows // Temporary sql rows

	rows, err = GetStatsByContainerIdStmt.Query(containerID)
	if err != nil {
		l.Error("GetStatsByContainerID:", err)
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		err = rows.Scan(&tmpStat.ContainerID, &tmpStat.Time, &tmpStat.SizeRootFs, &tmpStat.SizeRw, &tmpStat.SizeMemory, &tmpStat.Running)
		if err != nil {
			l.Error("GetStatsByContainerID:", err)
			return nil, err
		}
		stats = append(stats, tmpStat)
	}
	err = rows.Err()
	if err != nil {
		l.Error("GetStatsByContainerID:", err)
		return nil, err
	}

	return stats, nil
}
<commit_msg>Add ProbesID map to retrieve probes<commit_after>package core

import (
	"database/sql"
	"errors"

	_ "github.com/go-sql-driver/mysql"

	"../utils"
)

/*
	Container info
*/
type Container struct {
	ID         string
	ProbeID    int
	Hostname   string
	Image      string
	IPAddress  string
	MacAddress string
}

/*
	Container's stats
*/
type Stat struct {
	ContainerID string
	Time        int64
	SizeRootFs  uint64
	SizeRw      uint64
	SizeMemory  uint64
	Running     bool
}

/*
	SQL Variables
*/
var (
	// DB
	DB       *sql.DB
	ProbesID map[string]int

	// Prepared queries
	GetProbeIDStmt            *sql.Stmt
	InsertProbeStmt           *sql.Stmt
	InsertContainerStmt       *sql.Stmt
	DeleteContainerStmt       *sql.Stmt
	InsertStatStmt            *sql.Stmt
	DeleteStatStmt            *sql.Stmt
	GetContainerByIdStmt      *sql.Stmt
	GetStatsByContainerIdStmt *sql.Stmt
)

func InitSQL() {
	var err error // Error handling

	// Connect DB
	sqlc := DGConfig.DockerGuard.SQLServer // SQL Config
	DB, err = sql.Open("mysql", sqlc.User+":"+sqlc.Pass+"@tcp("+sqlc.IP+":"+utils.I2S(sqlc.Port)+")/"+sqlc.DB)
	if err != nil {
		l.Critical(err)
	}
	err = DB.Ping()
	if err != nil {
		l.Critical(err)
	}
	l.Verbose("Connected to SQL database")

	// Prepare DB queries
	GetProbeIDStmt, err = DB.Prepare("SELECT id FROM probes WHERE name=?")
	if err != nil {
		l.Critical("Can't create InsertContainerStmt:", err)
	}
	InsertProbeStmt, err = DB.Prepare("INSERT INTO probes VALUES (DEFAULT,?)")
	if err != nil {
		l.Critical("Can't create InsertContainerStmt:", err)
	}
	InsertContainerStmt, err = DB.Prepare("INSERT INTO containers VALUES (?,?,?,?,?,?)")
	if err != nil {
		l.Critical("Can't create InsertContainerStmt:", err)
	}
	DeleteContainerStmt, err = DB.Prepare("DELETE FROM containers WHERE id=?")
	if err != nil {
		l.Critical("Can't create DeleteContainerStmt:", err)
	}
	InsertStatStmt, err = DB.Prepare("INSERT INTO stats VALUES (?,?,?,?,?,?)")
	if err != nil {
		l.Critical("Can't create InsertStatStmt:", err)
	}
	DeleteStatStmt, err = DB.Prepare("DELETE FROM stats WHERE containerid=? AND time=?")
	if err != nil {
		l.Critical("Can't create DeleteStatStmt:", err)
	}
	GetContainerByIdStmt, err = DB.Prepare("SELECT * FROM containers WHERE id=?")
	if err != nil {
		l.Critical("Can't create GetContainerByIdStmt:", err)
	}
	GetStatsByContainerIdStmt, err = DB.Prepare("SELECT * FROM stats WHERE containerid=?")
	if err != nil {
		l.Critical("Can't create DeleteStatStmt:", err)
	}

	// Get probes ID
	ProbesID = make(map[string]int)
	for _, probe := range DGConfig.Probes {
		id, err := GetProbeID(probe.Name)
		if err != nil {
			l.Critical("Error GetProbeID ("+probe.Name+"):", err)
		}
		ProbesID[probe.Name] = id
	}
}

/*
	Get a probe ID from sql server
	If does not exist: create it and return the probe ID
*/
func GetProbeID(name string) (int, error) {
	var id int64  // Probe ID to return
	var err error // Error handling

	err = GetProbeIDStmt.QueryRow(name).Scan(&id)
	if err != nil {
		if err.Error() == "sql: no rows in result set" {
			result, err := InsertProbeStmt.Exec(name)
			if err != nil {
				l.Error("GetProbeID:", err)
				return 0, err
			}
			id, err = result.LastInsertId()
			if err != nil {
				l.Error("GetProbeID:", err)
				return 0, err
			}
		} else {
			l.Error("GetProbeID:", err)
			return 0, err
		}
	}

	return int(id), nil
}

/*
	Insert a container
*/
func (c *Container) Insert() error {
	var err error // Error handling

	_, err = InsertContainerStmt.Exec(c.ID, c.ProbeID, c.Hostname, c.Image, c.IPAddress, c.MacAddress)

	return err
}

/*
	Delete a container
*/
func (c *Container) Delete() error {
	var err error // Error handling

	_, err = DeleteContainerStmt.Exec(c.ID)

	return err
}

/*
	Get containers by $field
*/
func GetContainersBy(field string, value interface{}) ([]Container, error) {
	var containers []Container // Containers to return
	var rows *sql.Rows         // SQL Rows
	var err error              // Error handling

	// Protection against SQL injection
	var fieldExists bool = false
	for _, i := range []string{"id", "hostname", "image", "ip", "mac"} {
		if field == i {
			fieldExists = true
		}
	}
	if !fieldExists {
		l.Error("GetContainersBy: Field (" + field + ") does not exist.")
		return containers, errors.New("GetContainersBy: Field (" + field + ") does not exist.")
	}

	// Get containers
	rows, err = DB.Query("SELECT * FROM containers WHERE "+field+"=?", value)
	if err != nil {
		l.Error("GetContainersBy:", err)
		return containers, err
	}
	defer rows.Close()
	for rows.Next() {
		var tmpContainer Container

		if err = rows.Scan(&tmpContainer.ID,
			&tmpContainer.ProbeID,
			&tmpContainer.Hostname,
			&tmpContainer.Image,
			&tmpContainer.IPAddress,
			&tmpContainer.MacAddress); err != nil {
			l.Error("GetContainersBy:", err)
			return containers, err
		}

		containers = append(containers, tmpContainer)
	}
	if err = rows.Err(); err != nil {
		l.Error("GetContainersBy:", err)
		return containers, err
	}

	return containers, nil
}

/*
	Get a container by id
*/
func GetContainerById(id string) (Container, error) {
	var container Container // Container to return
	var err error           // Error handling

	err = GetContainerByIdStmt.QueryRow(id).Scan(&container.ID, &container.ProbeID, &container.Hostname, &container.Image, &container.IPAddress, &container.MacAddress)
	if err != nil {
		l.Error("GetContainerById:", err)
		return container, err
	}

	return container, nil
}

/*
	Insert a stat
*/
func (s *Stat) Insert() error {
	var err error // Error handling

	_, err = InsertStatStmt.Exec(s.ContainerID, s.Time, s.SizeRootFs, s.SizeRw, s.SizeMemory, s.Running)

	return err
}

/*
	Delete a stat
*/
func (s *Stat) Delete() error {
	var err error // Error handling

	_, err = DeleteStatStmt.Exec(s.ContainerID, s.Time)

	return err
}

/*
	Get stats by container id
*/
func GetStatsByContainerID(containerID string) ([]Stat, error) {
	var stats []Stat   // List of stats to return
	var err error      // Erro handling
	var tmpStat Stat   // Temporary stat
	var rows *sql.Rows // Temporary sql rows

	rows, err = GetStatsByContainerIdStmt.Query(containerID)
	if err != nil {
		l.Error("GetStatsByContainerID:", err)
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		err = rows.Scan(&tmpStat.ContainerID, &tmpStat.Time, &tmpStat.SizeRootFs, &tmpStat.SizeRw, &tmpStat.SizeMemory, &tmpStat.Running)
		if err != nil {
			l.Error("GetStatsByContainerID:", err)
			return nil, err
		}
		stats = append(stats, tmpStat)
	}
	err = rows.Err()
	if err != nil {
		l.Error("GetStatsByContainerID:", err)
		return nil, err
	}

	return stats, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package integration

import (
	"context"
	"fmt"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	prowjobv1 "k8s.io/test-infra/prow/apis/prowjobs/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	ctrlruntimeclient "sigs.k8s.io/controller-runtime/pkg/client"
)

func TestDeletePod(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name       string
		hasCR      bool
		wantDelete bool
	}{
		{
			"running-pod",
			true,
			false,
		},
		{
			"orphaned-pod",
			false,
			true,
		},
	}

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			clusterContext := getClusterContext()
			t.Logf("Creating client for cluster: %s", clusterContext)

			kubeClient, err := NewClients("", clusterContext)
			if err != nil {
				t.Fatalf("Failed creating clients for cluster %q: %v", clusterContext, err)
			}

			ctx := context.Background()

			prowjob := prowjobv1.ProwJob{
				ObjectMeta: v1.ObjectMeta{
					Annotations: map[string]string{
						"prow.k8s.io/job": tt.name,
					},
					Labels: map[string]string{
						"created-by-prow":  "true",
						"prow.k8s.io/type": "periodic",
						"name":             tt.name,
					},
					Name:      tt.name,
					Namespace: defaultNamespace,
				},
				Spec: prowjobv1.ProwJobSpec{
					Type:      prowjobv1.PeriodicJob,
					Namespace: testpodNamespace,
					Job:       tt.name,
				},
				Status: prowjobv1.ProwJobStatus{
					State: prowjobv1.TriggeredState,
				},
			}
			pod := corev1.Pod{
				ObjectMeta: v1.ObjectMeta{
					Name:      tt.name,
					Namespace: testpodNamespace,
					Labels: map[string]string{
						"name":            tt.name,
						"created-by-prow": "true",
					},
				},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{
						{
							Name:  "alpine",
							Image: "localhost:5000/alpine",
							Args: []string{
								"sleep",
								"1000000",
							},
						},
					},
				},
			}

			t.Cleanup(func() {
				kubeClient.Delete(ctx, &prowjob)
				kubeClient.Delete(ctx, &pod)
			})

			t.Logf("Creating prowjob: %s", tt.name)
			if err := kubeClient.Create(ctx, &prowjob); err != nil {
				t.Fatalf("Failed creating prowjob: %v", err)
			}
			t.Logf("Finished creating prowjob: %s", tt.name)

			// Create pod
			t.Logf("Creating pod: %s", tt.name)
			if err := kubeClient.Create(ctx, &pod); err != nil {
				t.Fatalf("Failed creating pod: %v", err)
			}
			t.Logf("Finished creating pod: %s", tt.name)

			// Make sure pod is running
			t.Logf("Make sure pod is running: %s", tt.name)
			if err = wait.Poll(time.Second, time.Minute, func() (bool, error) {
				p := &corev1.Pod{}
				if err := kubeClient.Get(ctx, client.ObjectKey{
					Name:      tt.name,
					Namespace: testpodNamespace,
				}, p); err != nil {
					return false, fmt.Errorf("Failed getting pod: %v", err)
				}
				return (p.Status.Phase == corev1.PodRunning), nil
			}); err != nil {
				t.Fatalf("Pod was not created successfully: %v", err)
			}
			t.Logf("Pod is running: %s", tt.name)

			// Delete prowjob if CR isn't supposed to have existed
			if !tt.hasCR {
				if err := kubeClient.Delete(ctx, &prowjob); err != nil {
					t.Fatalf("Failed deleting prowjob: %v", err)
				}
			}

			// Make sure pod is deleted, it'll take roughly 2 minutes
			// Don't care about the outcome, will check later
			t.Logf("Wait for sinker deleting pod or timeout in 2 minutes: %s", tt.name)
			var exist bool
			wait.Poll(time.Second, 2*time.Minute, func() (bool, error) {
				exist = false
				pods := &corev1.PodList{}
				err = kubeClient.List(ctx, pods, ctrlruntimeclient.InNamespace(testpodNamespace))
				if err != nil {
					return false, err
				}
				for _, p := range pods.Items {
					if p.Name == tt.name {
						exist = true
					}
				}
				return !exist, nil
			})
			// Check for the error of `List` call.
			if err != nil {
				t.Fatal(err)
			}
			t.Logf("Pod %s exist: %v", tt.name, exist)
			if want, got := tt.wantDelete, !exist; want != got {
				t.Fatalf("Want deleted: %v. Got deleted: %v", want, got)
			}
		})
	}
}
<commit_msg>Prow integration test: add more scenarios for sinker test<commit_after>/*
Copyright 2020 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package integration

import (
	"context"
	"fmt"
	"log"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	prowjobv1 "k8s.io/test-infra/prow/apis/prowjobs/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	ctrlruntimeclient "sigs.k8s.io/controller-runtime/pkg/client"
)

func TestDeletePod(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name             string
		notCreatedByProw bool
		jobStartTime     v1.Time
		jobFinishTime    v1.Time
		hasCR            bool
		wantJobDeleted   bool
		wantPodDeleted   bool
	}{
		{
			name:           "running-pod",
			jobStartTime:   v1.NewTime(time.Now().Add(-1 * time.Minute)),
			hasCR:          true,
			wantPodDeleted: false,
		},
		{
			name:           "orphaned-pod",
			jobStartTime:   v1.NewTime(time.Now().Add(-1 * time.Minute)),
			hasCR:          false,
			wantPodDeleted: true,
		},
		{
			name:           "ttl-deleted",
			jobStartTime:   v1.NewTime(time.Now().Add(-32 * time.Minute)),
			jobFinishTime:  v1.NewTime(time.Now().Add(-31 * time.Minute)),
			hasCR:          true,
			wantPodDeleted: true,
		},
		{
			name:           "max-prow-job-age",
			jobStartTime:   v1.NewTime(time.Now().Add(-49 * time.Hour)),
			jobFinishTime:  v1.NewTime(time.Now().Add(-48 * time.Hour)),
			hasCR:          true,
			wantJobDeleted: true,
			wantPodDeleted: true,
		},
		{
			name:             "orphaned-pod-not-prowpod",
			notCreatedByProw: true,
			wantPodDeleted:   false,
		},
	}

	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			resourceName := fmt.Sprintf("%s-%s", tt.name, RandomString(t))

			clusterContext := getClusterContext()
			t.Logf("Creating client for cluster: %s", clusterContext)

			kubeClient, err := NewClients("", clusterContext)
			if err != nil {
				t.Fatalf("Failed creating clients for cluster %q: %v", clusterContext, err)
			}

			ctx := context.Background()

			prowjob := prowjobv1.ProwJob{
				ObjectMeta: v1.ObjectMeta{
					Annotations: map[string]string{
						"prow.k8s.io/job": resourceName,
					},
					Labels: map[string]string{
						"created-by-prow":  "true",
						"prow.k8s.io/type": "periodic",
						"name":             resourceName,
					},
					Name:      resourceName,
					Namespace: defaultNamespace,
				},
				Spec: prowjobv1.ProwJobSpec{
					Type:      prowjobv1.PeriodicJob,
					Namespace: testpodNamespace,
					Job:       resourceName,
				},
				Status: prowjobv1.ProwJobStatus{
					State: prowjobv1.TriggeredState,
				},
			}

			pod := corev1.Pod{
				ObjectMeta: v1.ObjectMeta{
					Name:      resourceName,
					Namespace: testpodNamespace,
					Labels: map[string]string{
						"name": resourceName,
					},
				},
				Spec: corev1.PodSpec{
					Containers: []corev1.Container{
						{
							Name:  "alpine",
							Image: "localhost:5000/alpine",
							Args: []string{
								"sleep",
								"1000000",
							},
						},
					},
				},
			}
			if !tt.notCreatedByProw {
				pod.ObjectMeta.Labels["created-by-prow"] = "true"
			}

			t.Cleanup(func() {
				kubeClient.Delete(ctx, &prowjob)
				kubeClient.Delete(ctx, &pod)
			})

			if !tt.notCreatedByProw {
				t.Logf("Creating prowjob: %s", resourceName)
				if err := kubeClient.Create(ctx, &prowjob); err != nil {
					t.Fatalf("Failed creating prowjob: %v", err)
				}
				t.Logf("Finished creating prowjob: %s", resourceName)
				// Make sure pod is running
				p := &prowjobv1.ProwJob{}
				if err := kubeClient.Get(ctx, client.ObjectKey{
					Name:      resourceName,
					Namespace: defaultNamespace,
				}, p); err != nil {
					log.Fatalf("Prowjob %q not exist: %v", resourceName, err)
				}
			}

			// Create pod
			t.Logf("Creating pod: %s", resourceName)
			if err := kubeClient.Create(ctx, &pod); err != nil {
				t.Fatalf("Failed creating pod: %v", err)
			}
			t.Logf("Finished creating pod: %s", resourceName)

			// Make sure pod is running
			t.Logf("Make sure pod is running: %s", resourceName)
			if err = wait.Poll(time.Second, time.Minute, func() (bool, error) {
				p := &corev1.Pod{}
				if err := kubeClient.Get(ctx, client.ObjectKey{
					Name:      resourceName,
					Namespace: testpodNamespace,
				}, p); err != nil {
					return false, fmt.Errorf("Failed getting pod: %v", err)
				}
				return (p.Status.Phase == corev1.PodRunning), nil
			}); err != nil {
				t.Fatalf("Pod was not created successfully: %v", err)
			}
			t.Logf("Pod is running: %s", resourceName)

			if !tt.notCreatedByProw {
				// Patch prowjob to make it stale if needed
				prowjob.Status.StartTime = tt.jobStartTime
				prowjob.Status.CompletionTime = &tt.jobFinishTime
				if err := kubeClient.Update(ctx, &prowjob); err != nil {
					t.Fatalf("Failed updating prowjob %q: %v", resourceName, err)
				}
				// Delete prowjob if CR isn't supposed to have existed
				if !tt.hasCR {
					if err := kubeClient.Delete(ctx, &prowjob); err != nil {
						t.Fatalf("Failed deleting prowjob: %v", err)
					}
				}
			}

			// Make sure pod is deleted, it'll take roughly 2 minutes
			// Don't care about the outcome, will check later
			t.Logf("Wait for sinker deleting pod or timeout in 2 minutes: %s", resourceName)
			var exist bool
			wait.Poll(time.Second, 2*time.Minute, func() (bool, error) {
				exist = false
				pods := &corev1.PodList{}
				err = kubeClient.List(ctx, pods, ctrlruntimeclient.InNamespace(testpodNamespace))
				if err != nil {
					return false, err
				}
				for _, p := range pods.Items {
					if p.Name == resourceName {
						exist = true
					}
				}
				return !exist, nil
			})
			// Check for the error of `List` call.
			if err != nil {
				t.Fatal(err)
			}
			t.Logf("Pod %s exist: %v", resourceName, exist)
			if want, got := tt.wantPodDeleted, !exist; want != got {
				t.Fatalf("Want deleted: %v. Got deleted: %v", want, got)
			}

			// Check for prowjob deletion.
			{
				exist = false
				pjs := &prowjobv1.ProwJobList{}
				err = kubeClient.List(ctx, pjs, ctrlruntimeclient.InNamespace(defaultNamespace))
				if err != nil {
					log.Fatalf("Failed listing prowjobs: %v", err)
				}
				for _, pj := range pjs.Items {
					if pj.Name == resourceName {
						exist = true
					}
				}
				if tt.wantJobDeleted && exist {
					t.Fatalf("Pod exisentce mismatch. Want: %v, got: %v", tt.wantJobDeleted, exist)
				}
			}
		})
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The etcd-operator Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package version

var Version = "0.1.0"
<commit_msg>version: bump to 0.1.0+git<commit_after>// Copyright 2016 The etcd-operator Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package version

var Version = "0.1.0+git"
<|endoftext|>


<commit_before>package version

// VERSION ...
const VERSION = "2.4.2"
<commit_msg>Fix previous version bump. (#138)<commit_after>package version

// VERSION ...
const VERSION = "2.4.3"
<|endoftext|>


<commit_before>package version

// Version the library version number
const Version = "2.0.2"
<commit_msg>Bump to 2.1.0<commit_after>package version

// Version the library version number
const Version = "2.1.0"
<|endoftext|>


<commit_before>package termite

import (
	"log"
	"net"
	"os"
	"path/filepath"
	"rpc"
)

// Expose functionality for the local tool to use.
type LocalMaster struct {
	master *Master
	listener net.Listener
	server   *rpc.Server
}

func localStart(m *Master, sock string) {
	me := LocalMaster{
	master: m,
	}
	me.server = rpc.NewServer()
	me.server.Register(&me)
	me.start(sock)
}

func (me *LocalMaster) Run(req *WorkRequest, rep *WorkResponse) os.Error {
	if req.RanLocally {
		log.Println("Ran command locally:", req.Argv)
		return nil
	}

	return me.master.run(req, rep)
}

func (me *LocalMaster) Shutdown(input *int, output *int) os.Error {
	me.listener.Close()
	return nil
}

func (me *LocalMaster) RefreshAttributeCache(input *int, output *int) os.Error {
	log.Println("Refreshing attribute cache")
	me.master.refreshAttributeCache()
	log.Println("Refresh done")
	return nil
}

func (me *LocalMaster) InspectFile(req *AttrRequest, rep *AttrResponse) os.Error {
	return me.master.fileServer.GetAttr(req, rep)
}

func (me *LocalMaster) start(sock string) {
	absSock, err := filepath.Abs(sock)
	if err != nil {
		log.Fatal("abs", err)
	}

	fi, err := os.Stat(absSock)
	if fi != nil && fi.IsSocket() {
		conn, _ := net.Dial("unix", absSock)
		if conn != nil {
			conn.Close()
			log.Fatal("socket has someone listening: ", absSock)
		}
		// TODO - should check explicitly for the relevant error message.
		log.Println("removing dead socket", absSock)
		os.Remove(absSock)
	}

	me.listener, err = net.Listen("unix", absSock)
	defer os.Remove(absSock)
	if err != nil {
		log.Fatal("startLocalServer: ", err)
	}
	err = os.Chmod(absSock, 0700)
	if err != nil {
		log.Fatal("sock chmod", err)
	}

	writableRoot := ""
	writableRoot, err = filepath.EvalSymlinks(absSock)
	if err != nil {
		log.Fatal("EvalSymlinks", err)
	}
	writableRoot = filepath.Clean(writableRoot)
	writableRoot, _ = filepath.Split(writableRoot)
	writableRoot = filepath.Clean(writableRoot)

	me.master.writableRoot = writableRoot
	log.Println("accepting connections on", absSock)
	for {
		conn, err := me.listener.Accept()
		if err != nil {
			log.Fatal("listener.accept", err)
		}
		if !me.master.pending.Accept(conn) {
			go me.server.ServeConn(conn)
		}
	}
}
<commit_msg>Do master shutdown async, so we can send a reply.<commit_after>package termite

import (
	"log"
	"net"
	"os"
	"path/filepath"
	"rpc"
	"time"
)

// Expose functionality for the local tool to use.
type LocalMaster struct {
	master *Master
	listener net.Listener
	server   *rpc.Server
}

func localStart(m *Master, sock string) {
	me := LocalMaster{
	master: m,
	}
	me.server = rpc.NewServer()
	me.server.Register(&me)
	me.start(sock)
}

func (me *LocalMaster) Run(req *WorkRequest, rep *WorkResponse) os.Error {
	if req.RanLocally {
		log.Println("Ran command locally:", req.Argv)
		return nil
	}

	return me.master.run(req, rep)
}

func (me *LocalMaster) Shutdown(req *int, rep *int) os.Error {
	time.AfterFunc(1e8, func() {
	me.listener.Close()
	},)
	return nil
}

func (me *LocalMaster) RefreshAttributeCache(input *int, output *int) os.Error {
	log.Println("Refreshing attribute cache")
	me.master.refreshAttributeCache()
	log.Println("Refresh done")
	return nil
}

func (me *LocalMaster) InspectFile(req *AttrRequest, rep *AttrResponse) os.Error {
	return me.master.fileServer.GetAttr(req, rep)
}

func (me *LocalMaster) start(sock string) {
	absSock, err := filepath.Abs(sock)
	if err != nil {
		log.Fatal("abs", err)
	}

	fi, err := os.Stat(absSock)
	if fi != nil && fi.IsSocket() {
		conn, _ := net.Dial("unix", absSock)
		if conn != nil {
			conn.Close()
			log.Fatal("socket has someone listening: ", absSock)
		}
		// TODO - should check explicitly for the relevant error message.
		log.Println("removing dead socket", absSock)
		os.Remove(absSock)
	}

	me.listener, err = net.Listen("unix", absSock)
	defer os.Remove(absSock)
	if err != nil {
		log.Fatal("startLocalServer: ", err)
	}
	err = os.Chmod(absSock, 0700)
	if err != nil {
		log.Fatal("sock chmod", err)
	}

	writableRoot := ""
	writableRoot, err = filepath.EvalSymlinks(absSock)
	if err != nil {
		log.Fatal("EvalSymlinks", err)
	}
	writableRoot = filepath.Clean(writableRoot)
	writableRoot, _ = filepath.Split(writableRoot)
	writableRoot = filepath.Clean(writableRoot)

	me.master.writableRoot = writableRoot
	log.Println("accepting connections on", absSock)
	for {
		conn, err := me.listener.Accept()
		if err == os.EINVAL {
			break
		}
		if err != nil {
			log.Fatal("listener.accept: ", err)
		}
		if !me.master.pending.Accept(conn) {
			go me.server.ServeConn(conn)
		}
	}
}
<|endoftext|>


<commit_before>package mysql

import (
	. "github.com/GrappigPanda/notorious/database"
	"os"
	"testing"
	"time"
)

var DBCONN, _ = OpenConnection()

func TestOpenConn(t *testing.T) {
	dbConn, err := OpenConnection()
	if err != nil {
		t.Fatalf("%v", err)
	}
	InitDB(dbConn)
}

func TestAddWhitelistedTorrent(t *testing.T) {
	newTorrent := &White_Torrent{
		InfoHash:  "12345123451234512345",
		Name:      "Hello Kitty Island Adventure.exe",
		AddedBy:   "127.0.0.1",
		DateAdded: time.Now().Unix(),
	}

	if !newTorrent.AddWhitelistedTorrent(nil) {
		t.Fatalf("Failed to Add a whitelisted torrent")
	}
}

func TestGetWhitelistedTorrents(t *testing.T) {
	newTorrent := &White_Torrent{
		InfoHash:  "12345123GetWhitelistedTorrents",
		Name:      "Hello Kitty Island Adventure3.exe",
		AddedBy:   "127.0.0.1",
		DateAdded: time.Now().Unix(),
	}

	newTorrent2 := &White_Torrent{
		InfoHash:  "FFFFFFFFFFFFhitelistedTorrents",
		Name:      "Hello Kitty Island Adventure4.exe",
		AddedBy:   "127.0.0.1",
		DateAdded: time.Now().Unix(),
	}

	newTorrent.AddWhitelistedTorrent(nil)
	newTorrent2.AddWhitelistedTorrent(nil)

	_, err := GetWhitelistedTorrents(nil)
	if err != nil {
		t.Fatalf("Failed to get all whitelisted torrents: %v", err)
	}
}

func TestGetWhitelistedTorrent(t *testing.T) {
	newTorrent := &White_Torrent{
		InfoHash:  "12345123GetWhitelistedTorrent",
		Name:      "Hello Kitty Island Adventure2.exe",
		AddedBy:   "127.0.0.1",
		DateAdded: time.Now().Unix(),
	}

	newTorrent.AddWhitelistedTorrent(nil)

	retval, err := GetWhitelistedTorrent(nil, newTorrent.InfoHash)
	if err != nil {
		t.Fatalf("Failed to GetWhitelistedTorrent: %v", err)
	}

	if retval.InfoHash != newTorrent.InfoHash {
		t.Fatalf("Expected %v, got %v", retval.InfoHash,
			newTorrent.InfoHash)
	}
}

func TestUpdateStats(t *testing.T) {
	expectedReturn := &TrackerStats{
		Downloaded: 6,
		Uploaded:   21,
	}

	newStats := &TrackerStats{
		Downloaded: 1,
		Uploaded:   1,
	}
	DBCONN.Save(&newStats)

	UpdateStats(nil, 20, 5)

	retval := &TrackerStats{}
	DBCONN.First(&retval)
	if retval.Downloaded != expectedReturn.Downloaded {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Downloaded,
			retval.Downloaded)
	}

	if retval.Uploaded != expectedReturn.Uploaded {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Uploaded,
			retval.Uploaded)
	}
}

func TestUpdatePeerStats(t *testing.T) {
	expectedReturn := &Peer_Stats{
		Downloaded: 6,
		Uploaded:   21,
		Ip:         "127.0.0.1",
	}

	newPeer := &Peer_Stats{
		Downloaded: 1,
		Uploaded:   1,
		Ip:         "127.0.0.1",
	}

	DBCONN.Save(&newPeer)

	UpdatePeerStats(nil, 20, 5, "127.0.0.1")

	retval := &Peer_Stats{}
	DBCONN.First(&retval)

	if retval.Downloaded != expectedReturn.Downloaded {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Downloaded,
			retval.Downloaded)
	}

	if retval.Uploaded != expectedReturn.Uploaded {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Uploaded,
			retval.Uploaded)
	}

	if retval.Ip != expectedReturn.Ip {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Ip,
			retval.Ip)
	}
}

func TestMain(m *testing.M) {
	DBCONN.DropTableIfExists(
		&TrackerStats{},
		&Peer_Stats{},
		&Torrent{},
		&White_Torrent{},
	)
	os.Exit(m.Run())
}
<commit_msg>Issue #152<commit_after>package mysql

import (
	. "github.com/GrappigPanda/notorious/database"
	"os"
	"testing"
	"time"
)

var DBCONN, _ = OpenConnection()

func TestOpenConn(t *testing.T) {
	dbConn, err := OpenConnection()
	if err != nil {
		t.Fatalf("%v", err)
	}
	InitDB(dbConn)
}

func TestAddWhitelistedTorrent(t *testing.T) {
	newTorrent := &White_Torrent{
		InfoHash:  "12345123451234512345",
		Name:      "Hello Kitty Island Adventure.exe",
		AddedBy:   "127.0.0.1",
		DateAdded: time.Now().Unix(),
	}

	if !newTorrent.AddWhitelistedTorrent(DBCONN) {
		t.Fatalf("Failed to Add a whitelisted torrent")
	}
}

func TestGetWhitelistedTorrents(t *testing.T) {
	newTorrent := &White_Torrent{
		InfoHash:  "12345123GetWhitelistedTorrents",
		Name:      "Hello Kitty Island Adventure3.exe",
		AddedBy:   "127.0.0.1",
		DateAdded: time.Now().Unix(),
	}

	newTorrent2 := &White_Torrent{
		InfoHash:  "FFFFFFFFFFFFhitelistedTorrents",
		Name:      "Hello Kitty Island Adventure4.exe",
		AddedBy:   "127.0.0.1",
		DateAdded: time.Now().Unix(),
	}

	newTorrent.AddWhitelistedTorrent(DBCONN)
	newTorrent2.AddWhitelistedTorrent(DBCONN)

	_, err := GetWhitelistedTorrents(DBCONN)
	if err != nil {
		t.Fatalf("Failed to get all whitelisted torrents: %v", err)
	}
}

func TestGetWhitelistedTorrent(t *testing.T) {
	newTorrent := &White_Torrent{
		InfoHash:  "12345123GetWhitelistedTorrent",
		Name:      "Hello Kitty Island Adventure2.exe",
		AddedBy:   "127.0.0.1",
		DateAdded: time.Now().Unix(),
	}

	newTorrent.AddWhitelistedTorrent(DBCONN)

	retval, err := GetWhitelistedTorrent(nil, newTorrent.InfoHash)
	if err != nil {
		t.Fatalf("Failed to GetWhitelistedTorrent: %v", err)
	}

	if retval.InfoHash != newTorrent.InfoHash {
		t.Fatalf("Expected %v, got %v", retval.InfoHash,
			newTorrent.InfoHash)
	}
}

func TestUpdateStats(t *testing.T) {
	expectedReturn := &TrackerStats{
		Downloaded: 6,
		Uploaded:   21,
	}

	newStats := &TrackerStats{
		Downloaded: 1,
		Uploaded:   1,
	}
	DBCONN.Save(&newStats)

	UpdateStats(nil, 20, 5)

	retval := &TrackerStats{}
	DBCONN.First(&retval)
	if retval.Downloaded != expectedReturn.Downloaded {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Downloaded,
			retval.Downloaded)
	}

	if retval.Uploaded != expectedReturn.Uploaded {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Uploaded,
			retval.Uploaded)
	}
}

func TestUpdatePeerStats(t *testing.T) {
	expectedReturn := &Peer_Stats{
		Downloaded: 6,
		Uploaded:   21,
		Ip:         "127.0.0.1",
	}

	newPeer := &Peer_Stats{
		Downloaded: 1,
		Uploaded:   1,
		Ip:         "127.0.0.1",
	}

	DBCONN.Save(&newPeer)

	UpdatePeerStats(nil, 20, 5, "127.0.0.1")

	retval := &Peer_Stats{}
	DBCONN.First(&retval)

	if retval.Downloaded != expectedReturn.Downloaded {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Downloaded,
			retval.Downloaded)
	}

	if retval.Uploaded != expectedReturn.Uploaded {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Uploaded,
			retval.Uploaded)
	}

	if retval.Ip != expectedReturn.Ip {
		t.Fatalf("Expected %v, got %v",
			expectedReturn.Ip,
			retval.Ip)
	}
}

func TestMain(m *testing.M) {
	DBCONN.DropTableIfExists(
		&TrackerStats{},
		&Peer_Stats{},
		&Torrent{},
		&White_Torrent{},
	)
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package workqueue

import (
	"github.com/Matir/gobuster/client"
	"github.com/Matir/gobuster/logging"
	"github.com/Matir/gobuster/robots"
	"github.com/Matir/gobuster/util"
	"net/url"
	"sync"
)

// WorkQueue is a singleton that maintains the queue of work to be done.
// It reads from one input channel, verifies that the URL is in scope,
// queues it, then writes it to the work channel to be done.
type WorkQueue struct {
	// Elements to be worked on
	head *queueNode
	// End for cheap appends
	tail *queueNode
	// Number of items in queue, for stats
	queueLen int
	// Channel for URLs to be considered
	src chan *url.URL
	// Channel for URLs to be worked on
	dst chan *url.URL
	// filter to determine if a URL should be processed
	filter func(*url.URL) bool
	// channel to track done
	started chan bool
	// counter of work being done
	ctr WorkCounter
}

type queueNode struct {
	// next ptr
	next *queueNode
	// data
	data *url.URL
}

type QueueAddFunc func(...*url.URL)
type QueueAddCount func(int)
type QueueDoneFunc func(int)

func NewWorkQueue(queueSize int, scope []*url.URL, allowUpgrades bool) *WorkQueue {
	q := &WorkQueue{
		src:     make(chan *url.URL, queueSize),
		dst:     make(chan *url.URL, queueSize),
		filter:  makeScopeFunc(scope, allowUpgrades),
		started: make(chan bool, 1),
	}
	q.ctr.L = &sync.Mutex{}
	return q
}

func (q *WorkQueue) AddURLs(urls ...*url.URL) {
	q.ctr.Add(int64(len(urls)))
	for _, u := range urls {
		q.src <- u
	}
}

func (q *WorkQueue) InputFinished() {
	close(q.src)
}

func (q *WorkQueue) GetWorkChan() <-chan *url.URL {
	return q.dst
}

func (q *WorkQueue) Run() {
	defer close(q.dst)

	q.started <- true
	keepGoing := true
	for keepGoing {
		keepGoing = q.runStep()
	}
}

// Run a single step of the queue, returning true if we should continue
func (q *WorkQueue) runStep() bool {
	if q.head != nil {
		// If we have work to send, non-blocking read
		select {
		case u, ok := <-q.src:
			if !ok {
				for q.head != nil {
					q.dst <- q.pop()
				}
				return false
			}
			if q.filter(u) {
				q.push(u)
			} else {
				q.reject(u)
			}
		case q.dst <- q.peek():
			q.pop()
		}
	} else {
		// Blocking read and non-blocking send
		u, ok := <-q.src
		if !ok {
			return false
		}
		if !q.filter(u) {
			q.reject(u)
			return true
		}
		select {
		case q.dst <- u:
		default:
			q.push(u)
		}
	}
	return true
}

func (q *WorkQueue) RunInBackground() {
	go q.Run()
}

func (q *WorkQueue) WaitPipe() {
	<-q.started
	q.ctr.L.Lock()
	if q.ctr.todo == q.ctr.done {
		q.ctr.L.Unlock()
		return
	}
	q.ctr.Wait()
}

func (q *WorkQueue) GetAddFunc() QueueAddFunc {
	return func(urls ...*url.URL) {
		q.AddURLs(urls...)
	}
}

func (q *WorkQueue) GetAddCount() QueueAddCount {
	return func(c int) {
		q.ctr.Add(int64(c))
	}
}

func (q *WorkQueue) GetDoneFunc() QueueDoneFunc {
	return func(c int) {
		q.ctr.Done(int64(c))
	}
}

func (q *WorkQueue) SeedFromRobots(scope []*url.URL, clientFactory client.ClientFactory) {
	for _, scopeURL := range scope {
		robotsData, err := robots.GetRobotsForURL(scopeURL, clientFactory)
		if err != nil {
			logging.Logf(logging.LogWarning, "Unable to get robots.txt data: %s", err)
		} else {
			for _, path := range robotsData.GetAllPaths() {
				pathURL := *scopeURL
				pathURL.Path = path
				// Filter will handle if this is out of scope
				q.AddURLs(scopeURL.ResolveReference(&pathURL))
			}
		}
	}
}

func (q *WorkQueue) reject(u *url.URL) {
	logging.Logf(logging.LogDebug, "Workqueue rejecting %s", u.String())
	q.ctr.Done(1)
}

// Append URL to end of queue
func (q *WorkQueue) push(u *url.URL) {
	node := &queueNode{data: u}
	if q.tail != nil {
		q.tail.next = node
	} else {
		q.head = node
	}
	q.tail = node
	q.queueLen++
}

// Get URL from front of queue
func (q *WorkQueue) pop() *url.URL {
	node := q.head
	if node == nil {
		return nil
	}
	q.head = q.head.next
	if q.head == nil {
		q.tail = nil
	}
	q.queueLen--
	return node.data
}

// Get URL from front of queue without removal
func (q *WorkQueue) peek() *url.URL {
	if q.head != nil {
		return q.head.data
	}
	return nil
}

// Build a function to check if the target URL is in scope.
func makeScopeFunc(scope []*url.URL, allowUpgrades bool) func(*url.URL) bool {
	allowedScopes := make([]*url.URL, len(scope))
	copy(allowedScopes, scope)
	if allowUpgrades {
		for _, scopeURL := range scope {
			if scopeURL.Scheme == "http" {
				deref := *scopeURL
				clone := &deref // Can't find a way to do this in one statement
				clone.Scheme = "https"
				allowedScopes = append(allowedScopes, clone)
			}
		}
	}
	return func(target *url.URL) bool {
		for _, scopeURL := range scope {
			if util.URLIsSubpath(scopeURL, target) {
				return true
			}
		}
		return false
	}
}
<commit_msg>Fix bug allowing HTTP->HTTPS upgrade.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package workqueue

import (
	"github.com/Matir/gobuster/client"
	"github.com/Matir/gobuster/logging"
	"github.com/Matir/gobuster/robots"
	"github.com/Matir/gobuster/util"
	"net/url"
	"sync"
)

// WorkQueue is a singleton that maintains the queue of work to be done.
// It reads from one input channel, verifies that the URL is in scope,
// queues it, then writes it to the work channel to be done.
type WorkQueue struct {
	// Elements to be worked on
	head *queueNode
	// End for cheap appends
	tail *queueNode
	// Number of items in queue, for stats
	queueLen int
	// Channel for URLs to be considered
	src chan *url.URL
	// Channel for URLs to be worked on
	dst chan *url.URL
	// filter to determine if a URL should be processed
	filter func(*url.URL) bool
	// channel to track done
	started chan bool
	// counter of work being done
	ctr WorkCounter
}

type queueNode struct {
	// next ptr
	next *queueNode
	// data
	data *url.URL
}

type QueueAddFunc func(...*url.URL)
type QueueAddCount func(int)
type QueueDoneFunc func(int)

func NewWorkQueue(queueSize int, scope []*url.URL, allowUpgrades bool) *WorkQueue {
	q := &WorkQueue{
		src:     make(chan *url.URL, queueSize),
		dst:     make(chan *url.URL, queueSize),
		filter:  makeScopeFunc(scope, allowUpgrades),
		started: make(chan bool, 1),
	}
	q.ctr.L = &sync.Mutex{}
	return q
}

func (q *WorkQueue) AddURLs(urls ...*url.URL) {
	q.ctr.Add(int64(len(urls)))
	for _, u := range urls {
		q.src <- u
	}
}

func (q *WorkQueue) InputFinished() {
	close(q.src)
}

func (q *WorkQueue) GetWorkChan() <-chan *url.URL {
	return q.dst
}

func (q *WorkQueue) Run() {
	defer close(q.dst)

	q.started <- true
	keepGoing := true
	for keepGoing {
		keepGoing = q.runStep()
	}
}

// Run a single step of the queue, returning true if we should continue
func (q *WorkQueue) runStep() bool {
	if q.head != nil {
		// If we have work to send, non-blocking read
		select {
		case u, ok := <-q.src:
			if !ok {
				for q.head != nil {
					q.dst <- q.pop()
				}
				return false
			}
			if q.filter(u) {
				q.push(u)
			} else {
				q.reject(u)
			}
		case q.dst <- q.peek():
			q.pop()
		}
	} else {
		// Blocking read and non-blocking send
		u, ok := <-q.src
		if !ok {
			return false
		}
		if !q.filter(u) {
			q.reject(u)
			return true
		}
		select {
		case q.dst <- u:
		default:
			q.push(u)
		}
	}
	return true
}

func (q *WorkQueue) RunInBackground() {
	go q.Run()
}

func (q *WorkQueue) WaitPipe() {
	<-q.started
	q.ctr.L.Lock()
	if q.ctr.todo == q.ctr.done {
		q.ctr.L.Unlock()
		return
	}
	q.ctr.Wait()
}

func (q *WorkQueue) GetAddFunc() QueueAddFunc {
	return func(urls ...*url.URL) {
		q.AddURLs(urls...)
	}
}

func (q *WorkQueue) GetAddCount() QueueAddCount {
	return func(c int) {
		q.ctr.Add(int64(c))
	}
}

func (q *WorkQueue) GetDoneFunc() QueueDoneFunc {
	return func(c int) {
		q.ctr.Done(int64(c))
	}
}

func (q *WorkQueue) SeedFromRobots(scope []*url.URL, clientFactory client.ClientFactory) {
	for _, scopeURL := range scope {
		robotsData, err := robots.GetRobotsForURL(scopeURL, clientFactory)
		if err != nil {
			logging.Logf(logging.LogWarning, "Unable to get robots.txt data: %s", err)
		} else {
			for _, path := range robotsData.GetAllPaths() {
				pathURL := *scopeURL
				pathURL.Path = path
				// Filter will handle if this is out of scope
				q.AddURLs(scopeURL.ResolveReference(&pathURL))
			}
		}
	}
}

func (q *WorkQueue) reject(u *url.URL) {
	logging.Logf(logging.LogDebug, "Workqueue rejecting %s", u.String())
	q.ctr.Done(1)
}

// Append URL to end of queue
func (q *WorkQueue) push(u *url.URL) {
	node := &queueNode{data: u}
	if q.tail != nil {
		q.tail.next = node
	} else {
		q.head = node
	}
	q.tail = node
	q.queueLen++
}

// Get URL from front of queue
func (q *WorkQueue) pop() *url.URL {
	node := q.head
	if node == nil {
		return nil
	}
	q.head = q.head.next
	if q.head == nil {
		q.tail = nil
	}
	q.queueLen--
	return node.data
}

// Get URL from front of queue without removal
func (q *WorkQueue) peek() *url.URL {
	if q.head != nil {
		return q.head.data
	}
	return nil
}

// Build a function to check if the target URL is in scope.
func makeScopeFunc(scope []*url.URL, allowUpgrades bool) func(*url.URL) bool {
	allowedScopes := make([]*url.URL, len(scope))
	copy(allowedScopes, scope)
	if allowUpgrades {
		for _, scopeURL := range scope {
			if scopeURL.Scheme == "http" {
				deref := *scopeURL
				clone := &deref // Can't find a way to do this in one statement
				clone.Scheme = "https"
				allowedScopes = append(allowedScopes, clone)
			}
		}
	}
	return func(target *url.URL) bool {
		for _, scopeURL := range allowedScopes {
			if util.URLIsSubpath(scopeURL, target) {
				return true
			}
		}
		return false
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package services

import (
	"context"
	"sync"
	"time"

	"github.com/gravitational/teleport/lib/backend"
	"github.com/gravitational/teleport/lib/defaults"
	"github.com/gravitational/teleport/lib/utils"

	"github.com/gravitational/trace"
	log "github.com/sirupsen/logrus"
)

// NewProxyWatcherFunc creates a new instance of proxy watcher,
// used in tests
type NewProxyWatcherFunc func() (*ProxyWatcher, error)

// NewProxyWatcher returns a new instance of changeset
func NewProxyWatcher(cfg ProxyWatcherConfig) (*ProxyWatcher, error) {
	if err := cfg.CheckAndSetDefaults(); err != nil {
		return nil, trace.Wrap(err)
	}
	retry, err := utils.NewLinear(utils.LinearConfig{
		Step: cfg.RetryPeriod / 10,
		Max:  cfg.RetryPeriod,
	})
	if err != nil {
		return nil, trace.Wrap(err)
	}
	ctx, cancel := context.WithCancel(cfg.Context)
	p := &ProxyWatcher{
		ctx:                ctx,
		cancel:             cancel,
		RWMutex:            &sync.RWMutex{},
		retry:              retry,
		ProxyWatcherConfig: cfg,
		FieldLogger:        cfg.Entry,
	}
	go p.watchProxies()
	return p, nil
}

// ProxyWatcher is a resource built on top of the events,
// it monitors the additions and deletions to the proxies
type ProxyWatcher struct {
	*sync.RWMutex
	log.FieldLogger
	ProxyWatcherConfig

	// retry is used to manage backoff logic for watches
	retry utils.Retry

	// current is a list of the current servers
	// as reported by the watcher
	current []Server

	ctx    context.Context
	cancel context.CancelFunc
}

// ProxyWatcherConfig configures proxy watcher
type ProxyWatcherConfig struct {
	// Context is a parent context
	// controlling the lifecycle of the changeset
	Context context.Context
	// Component is a component used in logs
	Component string
	// RetryPeriod is a retry period on failed watchers
	RetryPeriod time.Duration
	// ReloadPeriod is a failed period on failed watches
	ReloadPeriod time.Duration
	// Client is used by changeset to monitor proxy updates
	Client ProxyWatcherClient
	// Entry is a logging entry
	Entry log.FieldLogger
	// ProxiesC is a channel that will be used
	// by the watcher to push updated list,
	// it will always receive a fresh list on the start
	// and the subsequent list of new values
	// whenever an addition or deletion to the list is detected
	ProxiesC chan []Server
}

// CheckAndSetDefaults checks parameters and sets default values
func (cfg *ProxyWatcherConfig) CheckAndSetDefaults() error {
	if cfg.Context == nil {
		return trace.BadParameter("missing parameter Context")
	}
	if cfg.Component == "" {
		return trace.BadParameter("missing parameter Component")
	}
	if cfg.Client == nil {
		return trace.BadParameter("missing parameter Client")
	}
	if cfg.ProxiesC == nil {
		return trace.BadParameter("missing parameter ProxiesC")
	}
	if cfg.Entry == nil {
		cfg.Entry = log.StandardLogger()
	}
	if cfg.RetryPeriod == 0 {
		cfg.RetryPeriod = defaults.HighResPollingPeriod
	}
	if cfg.ReloadPeriod == 0 {
		cfg.ReloadPeriod = defaults.LowResPollingPeriod
	}
	return nil
}

// GetCurrent returns a list of currently active proxies
func (p *ProxyWatcher) GetCurrent() []Server {
	p.RLock()
	defer p.RUnlock()
	return p.current
}

// setCurrent sets currently active proxy list
func (p *ProxyWatcher) setCurrent(proxies []Server) {
	p.Lock()
	defer p.Unlock()

	// Log the updated list of proxies. Useful for debugging to verify proxies
	// are being added and removed from the cluster.
	names := make([]string, 0, len(proxies))
	for _, proxy := range proxies {
		names = append(names, proxy.GetName())
	}
	p.Debugf("List of known proxies updated: %q.", names)

	p.current = proxies
}

// ProxyWatcherClient is used by changeset to fetch a list
// of proxies and subscribe to updates
type ProxyWatcherClient interface {
	ProxyGetter
	Events
}

// Done returns a channel that signals
// proxy watcher closure
func (p *ProxyWatcher) Done() <-chan struct{} {
	return p.ctx.Done()
}

// Close closes proxy watcher and cancels all the functions
func (p *ProxyWatcher) Close() error {
	p.cancel()
	return nil
}

// watchProxies watches new proxies added and removed to the cluster
// and when this happens, notifies all connected agents
// about the proxy set change via discovery requests
func (p *ProxyWatcher) watchProxies() {
	for {
		// Reload period is here to protect against
		// unknown cache going out of sync problems
		// that we did not predict.
		err := p.watch()
		if err != nil {
			p.Warningf("Re-init the watcher on error: %v.", trace.Unwrap(err))
		}
		p.Debugf("Reloading %v.", p.retry)
		select {
		case <-p.retry.After():
			p.retry.Inc()
		case <-p.ctx.Done():
			p.Debugf("Closed, returning from update loop.")
			return
		}
	}
}

// watch sets up the watch on proxies
func (p *ProxyWatcher) watch() error {
	watcher, err := p.Client.NewWatcher(p.Context, Watch{
		Name: p.Component,
		Kinds: []WatchKind{
			{
				Kind: KindProxy,
			},
		},
		MetricComponent: p.Component,
	})
	if err != nil {
		return trace.Wrap(err)
	}
	defer watcher.Close()
	reloadC := time.After(p.ReloadPeriod)
	// before fetch, make sure watcher is synced by receiving init event,
	// to avoid the scenario:
	// 1. Cache process:   w = NewWatcher()
	// 2. Cache process:   c.fetch()
	// 3. Backend process: addItem()
	// 4. Cache process:   <- w.Events()
	//
	// If there is a way that NewWatcher() on line 1 could
	// return without subscription established first,
	// Code line 3 could execute and line 4 could miss event,
	// wrapping up with out of sync replica.
	// To avoid this, before doing fetch,
	// cache process makes sure the connection is established
	// by receiving init event first.
	select {
	case <-watcher.Done():
		return trace.ConnectionProblem(watcher.Error(), "watcher is closed")
	case <-reloadC:
		p.Debugf("Triggering scheduled reload.")
		return nil
	case <-p.ctx.Done():
		return trace.ConnectionProblem(p.Context.Err(), "context is closing")
	case event := <-watcher.Events():
		if event.Type != backend.OpInit {
			return trace.BadParameter("expected init event, got %v instead", event.Type)
		}
	}
	proxies, err := p.Client.GetProxies()
	if err != nil {
		return trace.Wrap(err)
	}
	proxySet := make(map[string]Server, len(proxies))
	for i := range proxies {
		proxySet[proxies[i].GetName()] = proxies[i]
	}
	p.retry.Reset()
	p.setCurrent(proxies)
	select {
	case p.ProxiesC <- proxies:
	case <-p.ctx.Done():
		return trace.ConnectionProblem(p.Context.Err(), "context is closing")
	}
	for {
		select {
		case <-watcher.Done():
			return trace.ConnectionProblem(watcher.Error(), "watcher is closed")
		case <-reloadC:
			p.Debugf("Triggering scheduled reload.")
			return nil
		case <-p.ctx.Done():
			return trace.ConnectionProblem(p.Context.Err(), "context is closing")
		case event := <-watcher.Events():
			updated := p.processEvent(event, proxySet)
			if updated {
				proxies = setToList(proxySet)
				p.setCurrent(proxies)
				select {
				case p.ProxiesC <- proxies:
				case <-p.Context.Done():
					return trace.ConnectionProblem(p.Context.Err(), "context is closing")
				}
			}
		}
	}
}

// processEvent updates proxy map and returns true if the proxies list have been modified -
// the proxy has been either added or deleted
func (p *ProxyWatcher) processEvent(event Event, proxies map[string]Server) bool {
	if event.Resource.GetKind() != KindProxy {
		p.Warningf("Unexpected event: %v.")
		return false
	}
	switch event.Type {
	case backend.OpDelete:
		delete(proxies, event.Resource.GetName())
		// Always return true if the proxy has been deleted to trigger
		// broadcast cleanup.
		return true
	case backend.OpPut:
		_, existed := proxies[event.Resource.GetName()]
		resource, ok := event.Resource.(Server)
		if !ok {
			p.Warningf("unexpected type %T", event.Resource)
			return false
		}
		proxies[event.Resource.GetName()] = resource
		return !existed
	default:
		p.Warningf("Skipping unsupported event type %v.", event.Type)
		return false
	}
}

func setToList(in map[string]Server) []Server {
	out := make([]Server, 0, len(in))
	for key := range in {
		out = append(out, in[key])
	}
	return out
}
<commit_msg>re-init proxywatcher on empty proxy list<commit_after>/*
Copyright 2019 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package services

import (
	"context"
	"sync"
	"time"

	"github.com/gravitational/teleport/lib/backend"
	"github.com/gravitational/teleport/lib/defaults"
	"github.com/gravitational/teleport/lib/utils"

	"github.com/gravitational/trace"
	log "github.com/sirupsen/logrus"
)

// NewProxyWatcherFunc creates a new instance of proxy watcher,
// used in tests
type NewProxyWatcherFunc func() (*ProxyWatcher, error)

// NewProxyWatcher returns a new instance of changeset
func NewProxyWatcher(cfg ProxyWatcherConfig) (*ProxyWatcher, error) {
	if err := cfg.CheckAndSetDefaults(); err != nil {
		return nil, trace.Wrap(err)
	}
	retry, err := utils.NewLinear(utils.LinearConfig{
		Step: cfg.RetryPeriod / 10,
		Max:  cfg.RetryPeriod,
	})
	if err != nil {
		return nil, trace.Wrap(err)
	}
	ctx, cancel := context.WithCancel(cfg.Context)
	p := &ProxyWatcher{
		ctx:                ctx,
		cancel:             cancel,
		RWMutex:            &sync.RWMutex{},
		retry:              retry,
		ProxyWatcherConfig: cfg,
		FieldLogger:        cfg.Entry,
	}
	go p.watchProxies()
	return p, nil
}

// ProxyWatcher is a resource built on top of the events,
// it monitors the additions and deletions to the proxies
type ProxyWatcher struct {
	*sync.RWMutex
	log.FieldLogger
	ProxyWatcherConfig

	// retry is used to manage backoff logic for watches
	retry utils.Retry

	// current is a list of the current servers
	// as reported by the watcher
	current []Server

	ctx    context.Context
	cancel context.CancelFunc
}

// ProxyWatcherConfig configures proxy watcher
type ProxyWatcherConfig struct {
	// Context is a parent context
	// controlling the lifecycle of the changeset
	Context context.Context
	// Component is a component used in logs
	Component string
	// RetryPeriod is a retry period on failed watchers
	RetryPeriod time.Duration
	// ReloadPeriod is a failed period on failed watches
	ReloadPeriod time.Duration
	// Client is used by changeset to monitor proxy updates
	Client ProxyWatcherClient
	// Entry is a logging entry
	Entry log.FieldLogger
	// ProxiesC is a channel that will be used
	// by the watcher to push updated list,
	// it will always receive a fresh list on the start
	// and the subsequent list of new values
	// whenever an addition or deletion to the list is detected
	ProxiesC chan []Server
}

// CheckAndSetDefaults checks parameters and sets default values
func (cfg *ProxyWatcherConfig) CheckAndSetDefaults() error {
	if cfg.Context == nil {
		return trace.BadParameter("missing parameter Context")
	}
	if cfg.Component == "" {
		return trace.BadParameter("missing parameter Component")
	}
	if cfg.Client == nil {
		return trace.BadParameter("missing parameter Client")
	}
	if cfg.ProxiesC == nil {
		return trace.BadParameter("missing parameter ProxiesC")
	}
	if cfg.Entry == nil {
		cfg.Entry = log.StandardLogger()
	}
	if cfg.RetryPeriod == 0 {
		cfg.RetryPeriod = defaults.HighResPollingPeriod
	}
	if cfg.ReloadPeriod == 0 {
		cfg.ReloadPeriod = defaults.LowResPollingPeriod
	}
	return nil
}

// GetCurrent returns a list of currently active proxies
func (p *ProxyWatcher) GetCurrent() []Server {
	p.RLock()
	defer p.RUnlock()
	return p.current
}

// setCurrent sets currently active proxy list
func (p *ProxyWatcher) setCurrent(proxies []Server) {
	p.Lock()
	defer p.Unlock()

	// Log the updated list of proxies. Useful for debugging to verify proxies
	// are being added and removed from the cluster.
	names := make([]string, 0, len(proxies))
	for _, proxy := range proxies {
		names = append(names, proxy.GetName())
	}
	p.Debugf("List of known proxies updated: %q.", names)

	p.current = proxies
}

// ProxyWatcherClient is used by changeset to fetch a list
// of proxies and subscribe to updates
type ProxyWatcherClient interface {
	ProxyGetter
	Events
}

// Done returns a channel that signals
// proxy watcher closure
func (p *ProxyWatcher) Done() <-chan struct{} {
	return p.ctx.Done()
}

// Close closes proxy watcher and cancels all the functions
func (p *ProxyWatcher) Close() error {
	p.cancel()
	return nil
}

// watchProxies watches new proxies added and removed to the cluster
// and when this happens, notifies all connected agents
// about the proxy set change via discovery requests
func (p *ProxyWatcher) watchProxies() {
	for {
		// Reload period is here to protect against
		// unknown cache going out of sync problems
		// that we did not predict.
		err := p.watch()
		if err != nil {
			p.Warningf("Re-init the watcher on error: %v.", trace.Unwrap(err))
		}
		p.Debugf("Reloading %v.", p.retry)
		select {
		case <-p.retry.After():
			p.retry.Inc()
		case <-p.ctx.Done():
			p.Debugf("Closed, returning from update loop.")
			return
		}
	}
}

// watch sets up the watch on proxies
func (p *ProxyWatcher) watch() error {
	watcher, err := p.Client.NewWatcher(p.Context, Watch{
		Name: p.Component,
		Kinds: []WatchKind{
			{
				Kind: KindProxy,
			},
		},
		MetricComponent: p.Component,
	})
	if err != nil {
		return trace.Wrap(err)
	}
	defer watcher.Close()
	reloadC := time.After(p.ReloadPeriod)
	// before fetch, make sure watcher is synced by receiving init event,
	// to avoid the scenario:
	// 1. Cache process:   w = NewWatcher()
	// 2. Cache process:   c.fetch()
	// 3. Backend process: addItem()
	// 4. Cache process:   <- w.Events()
	//
	// If there is a way that NewWatcher() on line 1 could
	// return without subscription established first,
	// Code line 3 could execute and line 4 could miss event,
	// wrapping up with out of sync replica.
	// To avoid this, before doing fetch,
	// cache process makes sure the connection is established
	// by receiving init event first.
	select {
	case <-watcher.Done():
		return trace.ConnectionProblem(watcher.Error(), "watcher is closed")
	case <-reloadC:
		p.Debugf("Triggering scheduled reload.")
		return nil
	case <-p.ctx.Done():
		return trace.ConnectionProblem(p.Context.Err(), "context is closing")
	case event := <-watcher.Events():
		if event.Type != backend.OpInit {
			return trace.BadParameter("expected init event, got %v instead", event.Type)
		}
	}
	proxies, err := p.Client.GetProxies()
	if err != nil {
		return trace.Wrap(err)
	}
	if len(proxies) == 0 {
		// at least 1 proxy aught to exist; yield back to the outer
		// retry loop and try again.
		return trace.NotFound("empty proxy list")
	}
	proxySet := make(map[string]Server, len(proxies))
	for i := range proxies {
		proxySet[proxies[i].GetName()] = proxies[i]
	}
	p.retry.Reset()
	p.setCurrent(proxies)
	select {
	case p.ProxiesC <- proxies:
	case <-p.ctx.Done():
		return trace.ConnectionProblem(p.Context.Err(), "context is closing")
	}
	for {
		select {
		case <-watcher.Done():
			return trace.ConnectionProblem(watcher.Error(), "watcher is closed")
		case <-reloadC:
			p.Debugf("Triggering scheduled reload.")
			return nil
		case <-p.ctx.Done():
			return trace.ConnectionProblem(p.Context.Err(), "context is closing")
		case event := <-watcher.Events():
			updated := p.processEvent(event, proxySet)
			if updated {
				proxies = setToList(proxySet)
				p.setCurrent(proxies)
				select {
				case p.ProxiesC <- proxies:
				case <-p.Context.Done():
					return trace.ConnectionProblem(p.Context.Err(), "context is closing")
				}
			}
		}
	}
}

// processEvent updates proxy map and returns true if the proxies list have been modified -
// the proxy has been either added or deleted
func (p *ProxyWatcher) processEvent(event Event, proxies map[string]Server) bool {
	if event.Resource.GetKind() != KindProxy {
		p.Warningf("Unexpected event: %v.")
		return false
	}
	switch event.Type {
	case backend.OpDelete:
		delete(proxies, event.Resource.GetName())
		// Always return true if the proxy has been deleted to trigger
		// broadcast cleanup.
		return true
	case backend.OpPut:
		_, existed := proxies[event.Resource.GetName()]
		resource, ok := event.Resource.(Server)
		if !ok {
			p.Warningf("unexpected type %T", event.Resource)
			return false
		}
		proxies[event.Resource.GetName()] = resource
		return !existed
	default:
		p.Warningf("Skipping unsupported event type %v.", event.Type)
		return false
	}
}

func setToList(in map[string]Server) []Server {
	out := make([]Server, 0, len(in))
	for key := range in {
		out = append(out, in[key])
	}
	return out
}
<|endoftext|>