<commit_before>package helpers

import (
	"context"
	"database/sql"
	"errors"
	"fmt"

	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/dendrite/roomserver/auth"
	"github.com/matrix-org/dendrite/roomserver/state"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/storage/shared"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
)

// TODO: temporary package which has helper functions used by both internal/perform packages.
// Move these to a more sensible place.

func UpdateToInviteMembership(
	mu *shared.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent,
	roomVersion gomatrixserverlib.RoomVersion,
) ([]api.OutputEvent, error) {
	// We may have already sent the invite to the user, either because we are
	// reprocessing this event, or because the we received this invite from a
	// remote server via the federation invite API. In those cases we don't need
	// to send the event.
	needsSending, err := mu.SetToInvite(*add)
	if err != nil {
		return nil, err
	}
	if needsSending {
		// We notify the consumers using a special event even though we will
		// notify them about the change in current state as part of the normal
		// room event stream. This ensures that the consumers only have to
		// consider a single stream of events when determining whether a user
		// is invited, rather than having to combine multiple streams themselves.
		onie := api.OutputNewInviteEvent{
			Event:       add.Headered(roomVersion),
			RoomVersion: roomVersion,
		}
		updates = append(updates, api.OutputEvent{
			Type:           api.OutputTypeNewInviteEvent,
			NewInviteEvent: &onie,
		})
	}
	return updates, nil
}

func IsServerCurrentlyInRoom(ctx context.Context, db storage.Database, serverName gomatrixserverlib.ServerName, roomID string) (bool, error) {
	info, err := db.RoomInfo(ctx, roomID)
	if err != nil {
		return false, err
	}
	if info == nil {
		return false, fmt.Errorf("unknown room %s", roomID)
	}

	eventNIDs, err := db.GetMembershipEventNIDsForRoom(ctx, info.RoomNID, true, false)
	if err != nil {
		return false, err
	}

	events, err := db.Events(ctx, eventNIDs)
	if err != nil {
		return false, err
	}
	gmslEvents := make([]*gomatrixserverlib.Event, len(events))
	for i := range events {
		gmslEvents[i] = events[i].Event
	}
	return auth.IsAnyUserOnServerWithMembership(serverName, gmslEvents, gomatrixserverlib.Join), nil
}

func IsInvitePending(
	ctx context.Context, db storage.Database,
	roomID, userID string,
) (bool, string, string, error) {
	// Look up the room NID for the supplied room ID.
	info, err := db.RoomInfo(ctx, roomID)
	if err != nil {
		return false, "", "", fmt.Errorf("r.DB.RoomInfo: %w", err)
	}
	if info == nil {
		return false, "", "", fmt.Errorf("cannot get RoomInfo: unknown room ID %s", roomID)
	}

	// Look up the state key NID for the supplied user ID.
	targetUserNIDs, err := db.EventStateKeyNIDs(ctx, []string{userID})
	if err != nil {
		return false, "", "", fmt.Errorf("r.DB.EventStateKeyNIDs: %w", err)
	}
	targetUserNID, targetUserFound := targetUserNIDs[userID]
	if !targetUserFound {
		return false, "", "", fmt.Errorf("missing NID for user %q (%+v)", userID, targetUserNIDs)
	}

	// Let's see if we have an event active for the user in the room. If
	// we do then it will contain a server name that we can direct the
	// send_leave to.
	senderUserNIDs, eventIDs, err := db.GetInvitesForUser(ctx, info.RoomNID, targetUserNID)
	if err != nil {
		return false, "", "", fmt.Errorf("r.DB.GetInvitesForUser: %w", err)
	}
	if len(senderUserNIDs) == 0 {
		return false, "", "", nil
	}
	userNIDToEventID := make(map[types.EventStateKeyNID]string)
	for i, nid := range senderUserNIDs {
		userNIDToEventID[nid] = eventIDs[i]
	}

	// Look up the user ID from the NID.
	senderUsers, err := db.EventStateKeys(ctx, senderUserNIDs)
	if err != nil {
		return false, "", "", fmt.Errorf("r.DB.EventStateKeys: %w", err)
	}
	if len(senderUsers) == 0 {
		return false, "", "", fmt.Errorf("no senderUsers")
	}

	senderUser, senderUserFound := senderUsers[senderUserNIDs[0]]
	if !senderUserFound {
		return false, "", "", fmt.Errorf("missing user for NID %d (%+v)", senderUserNIDs[0], senderUsers)
	}

	return true, senderUser, userNIDToEventID[senderUserNIDs[0]], nil
}

// GetMembershipsAtState filters the state events to
// only keep the "m.room.member" events with a "join" membership. These events are returned.
// Returns an error if there was an issue fetching the events.
func GetMembershipsAtState(
	ctx context.Context, db storage.Database, stateEntries []types.StateEntry, joinedOnly bool,
) ([]types.Event, error) {

	var eventNIDs []types.EventNID
	for _, entry := range stateEntries {
		// Filter the events to retrieve to only keep the membership events
		if entry.EventTypeNID == types.MRoomMemberNID {
			eventNIDs = append(eventNIDs, entry.EventNID)
		}
	}

	// Get all of the events in this state
	stateEvents, err := db.Events(ctx, eventNIDs)
	if err != nil {
		return nil, err
	}

	if !joinedOnly {
		return stateEvents, nil
	}

	// Filter the events to only keep the "join" membership events
	var events []types.Event
	for _, event := range stateEvents {
		membership, err := event.Membership()
		if err != nil {
			return nil, err
		}

		if membership == gomatrixserverlib.Join {
			events = append(events, event)
		}
	}

	return events, nil
}

func StateBeforeEvent(ctx context.Context, db storage.Database, info types.RoomInfo, eventNID types.EventNID) ([]types.StateEntry, error) {
	roomState := state.NewStateResolution(db, info)
	// Lookup the event NID
	eIDs, err := db.EventIDs(ctx, []types.EventNID{eventNID})
	if err != nil {
		return nil, err
	}
	eventIDs := []string{eIDs[eventNID]}

	prevState, err := db.StateAtEventIDs(ctx, eventIDs)
	if err != nil {
		return nil, err
	}

	// Fetch the state as it was when this event was fired
	return roomState.LoadCombinedStateAfterEvents(ctx, prevState)
}

func LoadEvents(
	ctx context.Context, db storage.Database, eventNIDs []types.EventNID,
) ([]*gomatrixserverlib.Event, error) {
	stateEvents, err := db.Events(ctx, eventNIDs)
	if err != nil {
		return nil, err
	}

	result := make([]*gomatrixserverlib.Event, len(stateEvents))
	for i := range stateEvents {
		result[i] = stateEvents[i].Event
	}
	return result, nil
}

func LoadStateEvents(
	ctx context.Context, db storage.Database, stateEntries []types.StateEntry,
) ([]*gomatrixserverlib.Event, error) {
	eventNIDs := make([]types.EventNID, len(stateEntries))
	for i := range stateEntries {
		eventNIDs[i] = stateEntries[i].EventNID
	}
	return LoadEvents(ctx, db, eventNIDs)
}

func CheckServerAllowedToSeeEvent(
	ctx context.Context, db storage.Database, info types.RoomInfo, eventID string, serverName gomatrixserverlib.ServerName, isServerInRoom bool,
) (bool, error) {
	roomState := state.NewStateResolution(db, info)
	stateEntries, err := roomState.LoadStateAtEvent(ctx, eventID)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return false, nil
		}
		return false, err
	}

	// TODO: We probably want to make it so that we don't have to pull
	// out all the state if possible.
	stateAtEvent, err := LoadStateEvents(ctx, db, stateEntries)
	if err != nil {
		return false, err
	}

	return auth.IsServerAllowed(serverName, isServerInRoom, stateAtEvent), nil
}

// TODO: Remove this when we have tests to assert correctness of this function
// nolint:gocyclo
func ScanEventTree(
	ctx context.Context, db storage.Database, info types.RoomInfo, front []string, visited map[string]bool, limit int,
	serverName gomatrixserverlib.ServerName,
) ([]types.EventNID, error) {
	var resultNIDs []types.EventNID
	var err error
	var allowed bool
	var events []types.Event
	var next []string
	var pre string

	// TODO: add tests for this function to ensure it meets the contract that callers expect (and doc what that is supposed to be)
	// Currently, callers like PerformBackfill will call scanEventTree with a pre-populated `visited` map, assuming that by doing
	// so means that the events in that map will NOT be returned from this function. That is not currently true, resulting in
	// duplicate events being sent in response to /backfill requests.
	initialIgnoreList := make(map[string]bool, len(visited))
	for k, v := range visited {
		initialIgnoreList[k] = v
	}

	resultNIDs = make([]types.EventNID, 0, limit)

	var checkedServerInRoom bool
	var isServerInRoom bool

	// Loop through the event IDs to retrieve the requested events and go
	// through the whole tree (up to the provided limit) using the events'
	// "prev_event" key.
BFSLoop:
	for len(front) > 0 {
		// Prevent unnecessary allocations: reset the slice only when not empty.
		if len(next) > 0 {
			next = make([]string, 0)
		}
		// Retrieve the events to process from the database.
		events, err = db.EventsFromIDs(ctx, front)
		if err != nil {
			return resultNIDs, err
		}

		if !checkedServerInRoom && len(events) > 0 {
			// It's nasty that we have to extract the room ID from an event, but many federation requests
			// only talk in event IDs, no room IDs at all (!!!)
			ev := events[0]
			isServerInRoom, err = IsServerCurrentlyInRoom(ctx, db, serverName, ev.RoomID())
			if err != nil {
				util.GetLogger(ctx).WithError(err).Error("Failed to check if server is currently in room, assuming not.")
			}
			checkedServerInRoom = true
		}

		for _, ev := range events {
			// Break out of the loop if the provided limit is reached.
			if len(resultNIDs) == limit {
				break BFSLoop
			}

			if !initialIgnoreList[ev.EventID()] {
				// Update the list of events to retrieve.
				resultNIDs = append(resultNIDs, ev.EventNID)
			}
			// Loop through the event's parents.
			for _, pre = range ev.PrevEventIDs() {
				// Only add an event to the list of next events to process if it
				// hasn't been seen before.
				if !visited[pre] {
					visited[pre] = true
					allowed, err = CheckServerAllowedToSeeEvent(ctx, db, info, pre, serverName, isServerInRoom)
					if err != nil {
						util.GetLogger(ctx).WithField("server", serverName).WithField("event_id", pre).WithError(err).Error(
							"Error checking if allowed to see event",
						)
						// drop the error, as we will often error at the DB level if we don't have the prev_event itself. Let's
						// just return what we have.
						return resultNIDs, nil
					}

					// If the event hasn't been seen before and the HS
					// requesting to retrieve it is allowed to do so, add it to
					// the list of events to retrieve.
					if allowed {
						next = append(next, pre)
					} else {
						util.GetLogger(ctx).WithField("server", serverName).WithField("event_id", pre).Info("Not allowed to see event")
					}
				}
			}
		}
		// Repeat the same process with the parent events we just processed.
		front = next
	}

	return resultNIDs, err
}

func QueryLatestEventsAndState(
	ctx context.Context, db storage.Database,
	request *api.QueryLatestEventsAndStateRequest,
	response *api.QueryLatestEventsAndStateResponse,
) error {
	roomInfo, err := db.RoomInfo(ctx, request.RoomID)
	if err != nil {
		return err
	}
	if roomInfo == nil || roomInfo.IsStub {
		response.RoomExists = false
		return nil
	}

	roomState := state.NewStateResolution(db, *roomInfo)
	response.RoomExists = true
	response.RoomVersion = roomInfo.RoomVersion

	var currentStateSnapshotNID types.StateSnapshotNID
	response.LatestEvents, currentStateSnapshotNID, response.Depth, err =
		db.LatestEventIDs(ctx, roomInfo.RoomNID)
	if err != nil {
		return err
	}

	var stateEntries []types.StateEntry
	if len(request.StateToFetch) == 0 {
		// Look up all room state.
		stateEntries, err = roomState.LoadStateAtSnapshot(
			ctx, currentStateSnapshotNID,
		)
	} else {
		// Look up the current state for the requested tuples.
		stateEntries, err = roomState.LoadStateAtSnapshotForStringTuples(
			ctx, currentStateSnapshotNID, request.StateToFetch,
		)
	}
	if err != nil {
		return err
	}

	stateEvents, err := LoadStateEvents(ctx, db, stateEntries)
	if err != nil {
		return err
	}

	for _, event := range stateEvents {
		response.StateEvents = append(response.StateEvents, event.Headered(roomInfo.RoomVersion))
	}

	return nil
}
<commit_msg>Optimise CheckServerAllowedToSeeEvent (#1602)<commit_after>package helpers

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"

	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/dendrite/roomserver/auth"
	"github.com/matrix-org/dendrite/roomserver/state"
	"github.com/matrix-org/dendrite/roomserver/storage"
	"github.com/matrix-org/dendrite/roomserver/storage/shared"
	"github.com/matrix-org/dendrite/roomserver/types"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
)

// TODO: temporary package which has helper functions used by both internal/perform packages.
// Move these to a more sensible place.

func UpdateToInviteMembership(
	mu *shared.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent,
	roomVersion gomatrixserverlib.RoomVersion,
) ([]api.OutputEvent, error) {
	// We may have already sent the invite to the user, either because we are
	// reprocessing this event, or because the we received this invite from a
	// remote server via the federation invite API. In those cases we don't need
	// to send the event.
	needsSending, err := mu.SetToInvite(*add)
	if err != nil {
		return nil, err
	}
	if needsSending {
		// We notify the consumers using a special event even though we will
		// notify them about the change in current state as part of the normal
		// room event stream. This ensures that the consumers only have to
		// consider a single stream of events when determining whether a user
		// is invited, rather than having to combine multiple streams themselves.
		onie := api.OutputNewInviteEvent{
			Event:       add.Headered(roomVersion),
			RoomVersion: roomVersion,
		}
		updates = append(updates, api.OutputEvent{
			Type:           api.OutputTypeNewInviteEvent,
			NewInviteEvent: &onie,
		})
	}
	return updates, nil
}

func IsServerCurrentlyInRoom(ctx context.Context, db storage.Database, serverName gomatrixserverlib.ServerName, roomID string) (bool, error) {
	info, err := db.RoomInfo(ctx, roomID)
	if err != nil {
		return false, err
	}
	if info == nil {
		return false, fmt.Errorf("unknown room %s", roomID)
	}

	eventNIDs, err := db.GetMembershipEventNIDsForRoom(ctx, info.RoomNID, true, false)
	if err != nil {
		return false, err
	}

	events, err := db.Events(ctx, eventNIDs)
	if err != nil {
		return false, err
	}
	gmslEvents := make([]*gomatrixserverlib.Event, len(events))
	for i := range events {
		gmslEvents[i] = events[i].Event
	}
	return auth.IsAnyUserOnServerWithMembership(serverName, gmslEvents, gomatrixserverlib.Join), nil
}

func IsInvitePending(
	ctx context.Context, db storage.Database,
	roomID, userID string,
) (bool, string, string, error) {
	// Look up the room NID for the supplied room ID.
	info, err := db.RoomInfo(ctx, roomID)
	if err != nil {
		return false, "", "", fmt.Errorf("r.DB.RoomInfo: %w", err)
	}
	if info == nil {
		return false, "", "", fmt.Errorf("cannot get RoomInfo: unknown room ID %s", roomID)
	}

	// Look up the state key NID for the supplied user ID.
	targetUserNIDs, err := db.EventStateKeyNIDs(ctx, []string{userID})
	if err != nil {
		return false, "", "", fmt.Errorf("r.DB.EventStateKeyNIDs: %w", err)
	}
	targetUserNID, targetUserFound := targetUserNIDs[userID]
	if !targetUserFound {
		return false, "", "", fmt.Errorf("missing NID for user %q (%+v)", userID, targetUserNIDs)
	}

	// Let's see if we have an event active for the user in the room. If
	// we do then it will contain a server name that we can direct the
	// send_leave to.
	senderUserNIDs, eventIDs, err := db.GetInvitesForUser(ctx, info.RoomNID, targetUserNID)
	if err != nil {
		return false, "", "", fmt.Errorf("r.DB.GetInvitesForUser: %w", err)
	}
	if len(senderUserNIDs) == 0 {
		return false, "", "", nil
	}
	userNIDToEventID := make(map[types.EventStateKeyNID]string)
	for i, nid := range senderUserNIDs {
		userNIDToEventID[nid] = eventIDs[i]
	}

	// Look up the user ID from the NID.
	senderUsers, err := db.EventStateKeys(ctx, senderUserNIDs)
	if err != nil {
		return false, "", "", fmt.Errorf("r.DB.EventStateKeys: %w", err)
	}
	if len(senderUsers) == 0 {
		return false, "", "", fmt.Errorf("no senderUsers")
	}

	senderUser, senderUserFound := senderUsers[senderUserNIDs[0]]
	if !senderUserFound {
		return false, "", "", fmt.Errorf("missing user for NID %d (%+v)", senderUserNIDs[0], senderUsers)
	}

	return true, senderUser, userNIDToEventID[senderUserNIDs[0]], nil
}

// GetMembershipsAtState filters the state events to
// only keep the "m.room.member" events with a "join" membership. These events are returned.
// Returns an error if there was an issue fetching the events.
func GetMembershipsAtState(
	ctx context.Context, db storage.Database, stateEntries []types.StateEntry, joinedOnly bool,
) ([]types.Event, error) {

	var eventNIDs []types.EventNID
	for _, entry := range stateEntries {
		// Filter the events to retrieve to only keep the membership events
		if entry.EventTypeNID == types.MRoomMemberNID {
			eventNIDs = append(eventNIDs, entry.EventNID)
		}
	}

	// Get all of the events in this state
	stateEvents, err := db.Events(ctx, eventNIDs)
	if err != nil {
		return nil, err
	}

	if !joinedOnly {
		return stateEvents, nil
	}

	// Filter the events to only keep the "join" membership events
	var events []types.Event
	for _, event := range stateEvents {
		membership, err := event.Membership()
		if err != nil {
			return nil, err
		}

		if membership == gomatrixserverlib.Join {
			events = append(events, event)
		}
	}

	return events, nil
}

func StateBeforeEvent(ctx context.Context, db storage.Database, info types.RoomInfo, eventNID types.EventNID) ([]types.StateEntry, error) {
	roomState := state.NewStateResolution(db, info)
	// Lookup the event NID
	eIDs, err := db.EventIDs(ctx, []types.EventNID{eventNID})
	if err != nil {
		return nil, err
	}
	eventIDs := []string{eIDs[eventNID]}

	prevState, err := db.StateAtEventIDs(ctx, eventIDs)
	if err != nil {
		return nil, err
	}

	// Fetch the state as it was when this event was fired
	return roomState.LoadCombinedStateAfterEvents(ctx, prevState)
}

func LoadEvents(
	ctx context.Context, db storage.Database, eventNIDs []types.EventNID,
) ([]*gomatrixserverlib.Event, error) {
	stateEvents, err := db.Events(ctx, eventNIDs)
	if err != nil {
		return nil, err
	}

	result := make([]*gomatrixserverlib.Event, len(stateEvents))
	for i := range stateEvents {
		result[i] = stateEvents[i].Event
	}
	return result, nil
}

func LoadStateEvents(
	ctx context.Context, db storage.Database, stateEntries []types.StateEntry,
) ([]*gomatrixserverlib.Event, error) {
	eventNIDs := make([]types.EventNID, len(stateEntries))
	for i := range stateEntries {
		eventNIDs[i] = stateEntries[i].EventNID
	}
	return LoadEvents(ctx, db, eventNIDs)
}

func CheckServerAllowedToSeeEvent(
	ctx context.Context, db storage.Database, info types.RoomInfo, eventID string, serverName gomatrixserverlib.ServerName, isServerInRoom bool,
) (bool, error) {
	roomState := state.NewStateResolution(db, info)
	stateEntries, err := roomState.LoadStateAtEvent(ctx, eventID)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return false, nil
		}
		return false, fmt.Errorf("roomState.LoadStateAtEvent: %w", err)
	}

	// Extract all of the event state key NIDs from the room state.
	var stateKeyNIDs []types.EventStateKeyNID
	for _, entry := range stateEntries {
		stateKeyNIDs = append(stateKeyNIDs, entry.EventStateKeyNID)
	}

	// Then request those state key NIDs from the database.
	stateKeys, err := db.EventStateKeys(ctx, stateKeyNIDs)
	if err != nil {
		return false, fmt.Errorf("db.EventStateKeys: %w", err)
	}

	// If the event state key doesn't match the given servername
	// then we'll filter it out. This does preserve state keys that
	// are "" since these will contain history visibility etc.
	for nid, key := range stateKeys {
		if key != "" && !strings.HasSuffix(key, ":"+string(serverName)) {
			delete(stateKeys, nid)
		}
	}

	// Now filter through all of the state events for the room.
	// If the state key NID appears in the list of valid state
	// keys then we'll add it to the list of filtered entries.
	var filteredEntries []types.StateEntry
	for _, entry := range stateEntries {
		if _, ok := stateKeys[entry.EventStateKeyNID]; ok {
			filteredEntries = append(filteredEntries, entry)
		}
	}

	if len(filteredEntries) == 0 {
		return false, nil
	}

	stateAtEvent, err := LoadStateEvents(ctx, db, filteredEntries)
	if err != nil {
		return false, err
	}

	return auth.IsServerAllowed(serverName, isServerInRoom, stateAtEvent), nil
}

// TODO: Remove this when we have tests to assert correctness of this function
// nolint:gocyclo
func ScanEventTree(
	ctx context.Context, db storage.Database, info types.RoomInfo, front []string, visited map[string]bool, limit int,
	serverName gomatrixserverlib.ServerName,
) ([]types.EventNID, error) {
	var resultNIDs []types.EventNID
	var err error
	var allowed bool
	var events []types.Event
	var next []string
	var pre string

	// TODO: add tests for this function to ensure it meets the contract that callers expect (and doc what that is supposed to be)
	// Currently, callers like PerformBackfill will call scanEventTree with a pre-populated `visited` map, assuming that by doing
	// so means that the events in that map will NOT be returned from this function. That is not currently true, resulting in
	// duplicate events being sent in response to /backfill requests.
	initialIgnoreList := make(map[string]bool, len(visited))
	for k, v := range visited {
		initialIgnoreList[k] = v
	}

	resultNIDs = make([]types.EventNID, 0, limit)

	var checkedServerInRoom bool
	var isServerInRoom bool

	// Loop through the event IDs to retrieve the requested events and go
	// through the whole tree (up to the provided limit) using the events'
	// "prev_event" key.
BFSLoop:
	for len(front) > 0 {
		// Prevent unnecessary allocations: reset the slice only when not empty.
		if len(next) > 0 {
			next = make([]string, 0)
		}
		// Retrieve the events to process from the database.
		events, err = db.EventsFromIDs(ctx, front)
		if err != nil {
			return resultNIDs, err
		}

		if !checkedServerInRoom && len(events) > 0 {
			// It's nasty that we have to extract the room ID from an event, but many federation requests
			// only talk in event IDs, no room IDs at all (!!!)
			ev := events[0]
			isServerInRoom, err = IsServerCurrentlyInRoom(ctx, db, serverName, ev.RoomID())
			if err != nil {
				util.GetLogger(ctx).WithError(err).Error("Failed to check if server is currently in room, assuming not.")
			}
			checkedServerInRoom = true
		}

		for _, ev := range events {
			// Break out of the loop if the provided limit is reached.
			if len(resultNIDs) == limit {
				break BFSLoop
			}

			if !initialIgnoreList[ev.EventID()] {
				// Update the list of events to retrieve.
				resultNIDs = append(resultNIDs, ev.EventNID)
			}
			// Loop through the event's parents.
			for _, pre = range ev.PrevEventIDs() {
				// Only add an event to the list of next events to process if it
				// hasn't been seen before.
				if !visited[pre] {
					visited[pre] = true
					allowed, err = CheckServerAllowedToSeeEvent(ctx, db, info, pre, serverName, isServerInRoom)
					if err != nil {
						util.GetLogger(ctx).WithField("server", serverName).WithField("event_id", pre).WithError(err).Error(
							"Error checking if allowed to see event",
						)
						// drop the error, as we will often error at the DB level if we don't have the prev_event itself. Let's
						// just return what we have.
						return resultNIDs, nil
					}

					// If the event hasn't been seen before and the HS
					// requesting to retrieve it is allowed to do so, add it to
					// the list of events to retrieve.
					if allowed {
						next = append(next, pre)
					} else {
						util.GetLogger(ctx).WithField("server", serverName).WithField("event_id", pre).Info("Not allowed to see event")
					}
				}
			}
		}
		// Repeat the same process with the parent events we just processed.
		front = next
	}

	return resultNIDs, err
}

func QueryLatestEventsAndState(
	ctx context.Context, db storage.Database,
	request *api.QueryLatestEventsAndStateRequest,
	response *api.QueryLatestEventsAndStateResponse,
) error {
	roomInfo, err := db.RoomInfo(ctx, request.RoomID)
	if err != nil {
		return err
	}
	if roomInfo == nil || roomInfo.IsStub {
		response.RoomExists = false
		return nil
	}

	roomState := state.NewStateResolution(db, *roomInfo)
	response.RoomExists = true
	response.RoomVersion = roomInfo.RoomVersion

	var currentStateSnapshotNID types.StateSnapshotNID
	response.LatestEvents, currentStateSnapshotNID, response.Depth, err =
		db.LatestEventIDs(ctx, roomInfo.RoomNID)
	if err != nil {
		return err
	}

	var stateEntries []types.StateEntry
	if len(request.StateToFetch) == 0 {
		// Look up all room state.
		stateEntries, err = roomState.LoadStateAtSnapshot(
			ctx, currentStateSnapshotNID,
		)
	} else {
		// Look up the current state for the requested tuples.
		stateEntries, err = roomState.LoadStateAtSnapshotForStringTuples(
			ctx, currentStateSnapshotNID, request.StateToFetch,
		)
	}
	if err != nil {
		return err
	}

	stateEvents, err := LoadStateEvents(ctx, db, stateEntries)
	if err != nil {
		return err
	}

	for _, event := range stateEvents {
		response.StateEvents = append(response.StateEvents, event.Headered(roomInfo.RoomVersion))
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

package transport

import (
	"net/url"
	"path"
	"sort"
	"strings"
)

// URL contains all information needed to dial a remote node.
type URL struct {
	// TNC/modem/interface/network type.
	Scheme string

	// The host interface address.
	Host string

	// Host username (typically the local stations callsign) and password information.
	User *url.Userinfo

	// Target callsign.
	Target string

	// List of digipeaters ("path" between origin and target).
	Digis []string

	// List of query parameters.
	Params url.Values
}

// ParseURL parses a raw urlstring into an URL.
//
// scheme://(mycall(:password)@)(host)(/digi1/...)/targetcall
// Examples:
//   - winmor:///LA1B                       (Addresses LA1B on WINMOR).
//   - ax25://mycall@myaxport/LD5SK/LA1B-10 (Addresses LA1B-10 via LD5SK using AX.25-port "myaxport" and "MYCALL" as source callsign).
//
// The special query parameter host will override the host part of the path. (E.g. ax25:///LA1B?host=ax0 == ax25://ax0/LA1B).
func ParseURL(rawurl string) (*URL, error) {
	u, err := url.Parse(rawurl)
	if err != nil {
		return nil, err
	}

	// The digis and target should be all upper case
	u.Path = strings.ToUpper(u.Path)

	via, target := path.Split(u.Path)
	if len(target) < 3 {
		return nil, ErrInvalidTarget
	}

	url := &URL{
		Scheme: u.Scheme,
		Host:   u.Host,
		User:   u.User,
		Target: target,
		Params: u.Query(),
	}

	if str := url.Params.Get("host"); str != "" {
		url.Host = str
	}

	// Digis
	url.Digis = strings.Split(strings.Trim(via, "/"), "/")
	sort.Reverse(sort.StringSlice(url.Digis))
	if len(url.Digis) == 1 && url.Digis[0] == "" {
		url.Digis = []string{}
	}

	digisUnsupported := url.Scheme == "winmor" || url.Scheme == "ardop" || url.Scheme == "telnet"
	if len(url.Digis) > 0 && digisUnsupported {
		return url, ErrDigisUnsupported
	}

	return url, nil
}

// Set the URL.User's username (usually the source callsign).
func (u *URL) SetUser(call string) { u.User = url.User(call) }
<commit_msg>transport: Fix go vet issue<commit_after>// Copyright 2016 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

package transport

import (
	"net/url"
	"path"
	"sort"
	"strings"
)

// URL contains all information needed to dial a remote node.
type URL struct {
	// TNC/modem/interface/network type.
	Scheme string

	// The host interface address.
	Host string

	// Host username (typically the local stations callsign) and password information.
	User *url.Userinfo

	// Target callsign.
	Target string

	// List of digipeaters ("path" between origin and target).
	Digis []string

	// List of query parameters.
	Params url.Values
}

// ParseURL parses a raw urlstring into an URL.
//
// scheme://(mycall(:password)@)(host)(/digi1/...)/targetcall
// Examples:
//   - winmor:///LA1B                       (Addresses LA1B on WINMOR).
//   - ax25://mycall@myaxport/LD5SK/LA1B-10 (Addresses LA1B-10 via LD5SK using AX.25-port "myaxport" and "MYCALL" as source callsign).
//
// The special query parameter host will override the host part of the path. (E.g. ax25:///LA1B?host=ax0 == ax25://ax0/LA1B).
func ParseURL(rawurl string) (*URL, error) {
	u, err := url.Parse(rawurl)
	if err != nil {
		return nil, err
	}

	// The digis and target should be all upper case
	u.Path = strings.ToUpper(u.Path)

	via, target := path.Split(u.Path)
	if len(target) < 3 {
		return nil, ErrInvalidTarget
	}

	url := &URL{
		Scheme: u.Scheme,
		Host:   u.Host,
		User:   u.User,
		Target: target,
		Params: u.Query(),
	}

	if str := url.Params.Get("host"); str != "" {
		url.Host = str
	}

	// Digis
	url.Digis = strings.Split(strings.Trim(via, "/"), "/")
	_ = sort.Reverse(sort.StringSlice(url.Digis))
	if len(url.Digis) == 1 && url.Digis[0] == "" {
		url.Digis = []string{}
	}

	digisUnsupported := url.Scheme == "winmor" || url.Scheme == "ardop" || url.Scheme == "telnet"
	if len(url.Digis) > 0 && digisUnsupported {
		return url, ErrDigisUnsupported
	}

	return url, nil
}

// Set the URL.User's username (usually the source callsign).
func (u *URL) SetUser(call string) { u.User = url.User(call) }
<|endoftext|>


<commit_before>// Copyright 2015 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package proto3_proto

import (
	"bytes"
	"encoding/gob"
	"log"
	"os"
	"reflect"
	"testing"
)

type tEchoService struct {
	private int
}

func (p *tEchoService) Echo(in *Message, out *Message) error {
	var buf bytes.Buffer
	if err := gob.NewEncoder(&buf).Encode(in); err != nil {
		return err
	}
	if err := gob.NewDecoder(bytes.NewBuffer(buf.Bytes())).Decode(out); err != nil {
		return err
	}
	return nil
}

func TestMain(m *testing.M) {
	go func() {
		if err := ListenAndServeEchoService("tcp", "127.0.0.1:3000", new(tEchoService)); err != nil {
			log.Fatal(err)
		}
	}()
	os.Exit(m.Run())
}

func TestEchoService(t *testing.T) {
	c, err := DialEchoService("tcp", "127.0.0.1:3000")
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	in := Message{
		Name:         "github.com/chai2010/protorpc",
		Hilarity:     Message_PUNS,
		HeightInCm:   13,
		Data:         []byte("bin data"),
		ResultCount:  2<<35 + 1,
		TrueScotsman: true,
		Score:        3.14,
		Key:          []uint64{1, 1001},
		Nested:       &Nested{Bunny: "{{Bunny}}"},
		Terrain: map[string]*Nested{
			"A": &Nested{Bunny: "{{A}}"},
			"B": &Nested{Bunny: "{{B}}"},
		},
	}

	out, err := c.Echo(&in)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(&in, out) {
		t.Fatalf("not euqal, got = %v\n", &out)
	}
}
<commit_msg>test: start service in init func<commit_after>// Copyright 2015 <chaishushan{AT}gmail.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package proto3_proto

import (
	"bytes"
	"encoding/gob"
	"log"
	"os"
	"reflect"
	"testing"
)

type tEchoService struct {
	private int
}

func (p *tEchoService) Echo(in *Message, out *Message) error {
	var buf bytes.Buffer
	if err := gob.NewEncoder(&buf).Encode(in); err != nil {
		return err
	}
	if err := gob.NewDecoder(bytes.NewBuffer(buf.Bytes())).Decode(out); err != nil {
		return err
	}
	return nil
}

func init() {
	go func() {
		if err := ListenAndServeEchoService("tcp", "127.0.0.1:3000", new(tEchoService)); err != nil {
			log.Fatal(err)
		}
	}()
}

func TestMain(m *testing.M) {
	os.Exit(m.Run())
}

func TestEchoService(t *testing.T) {
	c, err := DialEchoService("tcp", "127.0.0.1:3000")
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()

	in := Message{
		Name:         "github.com/chai2010/protorpc",
		Hilarity:     Message_PUNS,
		HeightInCm:   13,
		Data:         []byte("bin data"),
		ResultCount:  2<<35 + 1,
		TrueScotsman: true,
		Score:        3.14,
		Key:          []uint64{1, 1001},
		Nested:       &Nested{Bunny: "{{Bunny}}"},
		Terrain: map[string]*Nested{
			"A": &Nested{Bunny: "{{A}}"},
			"B": &Nested{Bunny: "{{B}}"},
		},
	}

	out, err := c.Echo(&in)
	if err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(&in, out) {
		t.Fatalf("not euqal, got = %v\n", &out)
	}
}
<|endoftext|>


<commit_before>package syscallcompat

import (
	"path/filepath"
	"strings"
	"syscall"

	"github.com/rfjakob/gocryptfs/internal/tlog"
)

// OpenDirNofollow opens the dir at "relPath" in a way that is secure against
// symlink attacks. Symlinks that are part of "relPath" are never followed.
// This function is implemented by walking the directory tree, starting at
// "baseDir", using the Openat syscall with the O_NOFOLLOW flag.
// Symlinks that are part of the "baseDir" path are followed.
func OpenDirNofollow(baseDir string, relPath string) (fd int, err error) {
	if !filepath.IsAbs(baseDir) {
		tlog.Warn.Printf("BUG: OpenDirNofollow called with relative baseDir=%q", baseDir)
		return -1, syscall.EINVAL
	}
	if filepath.IsAbs(relPath) {
		tlog.Warn.Printf("BUG: OpenDirNofollow called with absolute relPath=%q", relPath)
		return -1, syscall.EINVAL
	}
	// Open the base dir (following symlinks)
	// TODO: should this use syscallcompat.O_PATH?
	dirfd, err := syscall.Open(baseDir, syscall.O_RDONLY|syscall.O_DIRECTORY, 0)
	if err != nil {
		return -1, err
	}
	// Caller wanted to open baseDir itself?
	if relPath == "" {
		return dirfd, nil
	}
	// Split the path into components
	parts := strings.Split(relPath, "/")
	// Walk the directory tree
	var dirfd2 int
	for _, name := range parts {
		dirfd2, err = Openat(dirfd, name, syscall.O_RDONLY|syscall.O_NOFOLLOW|syscall.O_DIRECTORY|O_PATH, 0)
		syscall.Close(dirfd)
		if err != nil {
			return -1, err
		}
		dirfd = dirfd2
	}
	// Return fd to final directory
	return dirfd, nil
}
<commit_msg>syscallcompat: Use O_PATH to open base directory.<commit_after>package syscallcompat

import (
	"path/filepath"
	"strings"
	"syscall"

	"github.com/rfjakob/gocryptfs/internal/tlog"
)

// OpenDirNofollow opens the dir at "relPath" in a way that is secure against
// symlink attacks. Symlinks that are part of "relPath" are never followed.
// This function is implemented by walking the directory tree, starting at
// "baseDir", using the Openat syscall with the O_NOFOLLOW flag.
// Symlinks that are part of the "baseDir" path are followed.
func OpenDirNofollow(baseDir string, relPath string) (fd int, err error) {
	if !filepath.IsAbs(baseDir) {
		tlog.Warn.Printf("BUG: OpenDirNofollow called with relative baseDir=%q", baseDir)
		return -1, syscall.EINVAL
	}
	if filepath.IsAbs(relPath) {
		tlog.Warn.Printf("BUG: OpenDirNofollow called with absolute relPath=%q", relPath)
		return -1, syscall.EINVAL
	}
	// Open the base dir (following symlinks)
	dirfd, err := syscall.Open(baseDir, syscall.O_DIRECTORY|O_PATH, 0)
	if err != nil {
		return -1, err
	}
	// Caller wanted to open baseDir itself?
	if relPath == "" {
		return dirfd, nil
	}
	// Split the path into components
	parts := strings.Split(relPath, "/")
	// Walk the directory tree
	var dirfd2 int
	for _, name := range parts {
		dirfd2, err = Openat(dirfd, name, syscall.O_NOFOLLOW|syscall.O_DIRECTORY|O_PATH, 0)
		syscall.Close(dirfd)
		if err != nil {
			return -1, err
		}
		dirfd = dirfd2
	}
	// Return fd to final directory
	return dirfd, nil
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 pkgbuilder

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
	"text/template"

	kptfileutil "github.com/GoogleContainerTools/kpt/pkg/kptfile"
	"github.com/stretchr/testify/assert"
	"sigs.k8s.io/kustomize/kyaml/yaml"
)

var (
	deploymentResourceManifest = `
apiVersion: apps/v1
kind: Deployment
metadata:
  namespace: myspace
  name: mysql-deployment
spec:
  replicas: 3
  foo: bar
  template:
    spec:
      containers:
      - name: mysql
        image: mysql:1.7.9
`

	configMapResourceManifest = `
apiVersion: v1
kind: ConfigMap
metadata:
  name: configmap
data:
  foo: bar
`
)

var (
	DeploymentResource = "deployment"
	ConfigMapResource  = "configmap"
	resources          = map[string]resourceInfo{
		DeploymentResource: {
			filename: "deployment.yaml",
			manifest: deploymentResourceManifest,
		},
		ConfigMapResource: {
			filename: "configmap.yaml",
			manifest: configMapResourceManifest,
		},
	}
)

type resourceInfo struct {
	filename string
	manifest string
}

type resourceInfoWithSetters struct {
	resourceInfo resourceInfo
	setterRefs   []SetterRef
	mutators     []yaml.Filter
}

// Pkg represents a package that can be created on the file system
// by using the Build function
type Pkg struct {
	Name string

	Kptfile *Kptfile

	resources []resourceInfoWithSetters

	subPkgs []*Pkg
}

func NewKptfile() *Kptfile {
	return &Kptfile{}
}

// Kptfile represents the Kptfile of a package.
type Kptfile struct {
	Setters []Setter
	Repo    string
	Ref     string
}

// WithUpstream adds information about the upstream information to the Kptfile.
// The upstream section of the Kptfile is only added if this information is
// provided.
func (k *Kptfile) WithUpstream(repo, ref string) *Kptfile {
	k.Repo = repo
	k.Ref = ref
	return k
}

// WithSetters adds information about the setters for a Kptfile.
func (k *Kptfile) WithSetters(setters ...Setter) *Kptfile {
	k.Setters = setters
	return k
}

// Setter contains the properties required for adding a setter to the
// Kptfile.
type Setter struct {
	Name  string
	Value string
	IsSet bool
}

// NewSetter creates a new setter that is not marked as set
func NewSetter(name, value string) Setter {
	return Setter{
		Name:  name,
		Value: value,
	}
}

// NewSetSetter creates a new setter that is marked as set.
func NewSetSetter(name, value string) Setter {
	return Setter{
		Name:  name,
		Value: value,
		IsSet: true,
	}
}

// SetterRef specifies the information for creating a new reference to
// a setter in a resource.
type SetterRef struct {
	Path []string
	Name string
}

// NewSetterRef creates a new setterRef with the given name and path.
func NewSetterRef(name string, path ...string) SetterRef {
	return SetterRef{
		Path: path,
		Name: name,
	}
}

// NewPackage creates a new package for testing.
func NewPackage(name string) *Pkg {
	return &Pkg{
		Name: name,
	}
}

// WithKptfile configures the current package to have a Kptfile. Only
// zero or one Kptfiles are accepted.
func (p *Pkg) WithKptfile(kf ...*Kptfile) *Pkg {
	if len(kf) > 1 {
		panic("only 0 or 1 Kptfiles are allowed")
	}
	if len(kf) == 0 {
		p.Kptfile = NewKptfile()
	} else {
		p.Kptfile = kf[0]
	}
	return p
}

// WithResource configures the package to include the provided resource
func (p *Pkg) WithResource(resourceName string, mutators ...yaml.Filter) *Pkg {
	resourceInfo, ok := resources[resourceName]
	if !ok {
		panic(fmt.Errorf("unknown resource %s", resourceName))
	}
	p.resources = append(p.resources, resourceInfoWithSetters{
		resourceInfo: resourceInfo,
		setterRefs:   []SetterRef{},
		mutators:     mutators,
	})
	return p
}

// WithResourceAndSetters configures the package to have the provided resource.
// It also allows for specifying setterRefs for the resource and a set of
// mutators that will update the content of the resource.
func (p *Pkg) WithResourceAndSetters(resourceName string, setterRefs []SetterRef, mutators ...yaml.Filter) *Pkg {
	resourceInfo, ok := resources[resourceName]
	if !ok {
		panic(fmt.Errorf("unknown resource %s", resourceName))
	}
	p.resources = append(p.resources, resourceInfoWithSetters{
		resourceInfo: resourceInfo,
		setterRefs:   setterRefs,
		mutators:     mutators,
	})
	return p
}

// WithSubPackages adds the provided packages as subpackages to the current
// package
func (p *Pkg) WithSubPackages(ps ...*Pkg) *Pkg {
	p.subPkgs = append(p.subPkgs, ps...)
	return p
}

// Build outputs the current data structure as a set of (nested) package
// in the provided path.
func (p *Pkg) Build(path string) error {
	_, err := os.Stat(path)
	if err != nil {
		return err
	}
	return buildRecursive(path, p)
}

func buildRecursive(path string, pkg *Pkg) error {
	pkgPath := filepath.Join(path, pkg.Name)
	err := os.Mkdir(pkgPath, 0700)
	if err != nil {
		return err
	}

	if pkg.Kptfile != nil {
		content := buildKptfile(pkg)

		err := ioutil.WriteFile(filepath.Join(pkgPath, kptfileutil.KptFileName),
			[]byte(content), 0600)
		if err != nil {
			return err
		}
	}

	for _, ri := range pkg.resources {
		m := ri.resourceInfo.manifest
		r := yaml.MustParse(m)
		for _, setterRef := range ri.setterRefs {
			n, err := r.Pipe(yaml.PathGetter{
				Path: setterRef.Path,
			})
			if err != nil {
				return err
			}
			n.YNode().LineComment = fmt.Sprintf(`{"$openapi":"%s"}`, setterRef.Name)
		}

		for _, m := range ri.mutators {
			if err := r.PipeE(m); err != nil {
				return err
			}
		}

		filePath := filepath.Join(pkgPath, ri.resourceInfo.filename)
		err = ioutil.WriteFile(filePath, []byte(r.MustString()), 0600)
		if err != nil {
			return err
		}
	}

	for i := range pkg.subPkgs {
		subPkg := pkg.subPkgs[i]
		err = buildRecursive(pkgPath, subPkg)
		if err != nil {
			return err
		}
	}

	return nil
}

var kptfileTemplate = `
apiVersion: kpt.dev/v1alpha1
kind: Kptfile
metadata:
  name: {{.Name}}
{{- if gt (len .Kptfile.Setters) 0 }}
openAPI:
  definitions:
{{- range .Kptfile.Setters }}
    io.k8s.cli.setters.{{.Name}}:
      x-k8s-cli:
        setter:
          name: {{.Name}}
          value: {{.Value}}
{{- if eq .IsSet true }}
          isSet: true
{{- end }}
{{- end }}
{{- end }}
{{- if gt (len .Kptfile.Repo) 0 }}
upstream:
  type: git
  git:
    ref: {{.Kptfile.Ref}}
    repo: {{.Kptfile.Repo}}
{{- end }}
`

func buildKptfile(pkg *Pkg) string {
	tmpl, err := template.New("test").Parse(kptfileTemplate)
	if err != nil {
		panic(err)
	}
	var buf bytes.Buffer
	err = tmpl.Execute(&buf, pkg)
	if err != nil {
		panic(err)
	}
	result := buf.String()
	return result
}

func ExpandPkg(t *testing.T, pkg *Pkg) string {
	if pkg.Name == "" {
		pkg.Name = "base"
	}
	dir, err := ioutil.TempDir("", "test-kpt-builder-")
	if !assert.NoError(t, err) {
		t.FailNow()
	}
	err = pkg.Build(dir)
	if !assert.NoError(t, err) {
		t.FailNow()
	}
	return filepath.Join(dir, pkg.Name)
}
<commit_msg>Support raw files in package builder<commit_after>// Copyright 2020 Google LLC
//
// 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 pkgbuilder

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
	"text/template"

	kptfileutil "github.com/GoogleContainerTools/kpt/pkg/kptfile"
	"github.com/stretchr/testify/assert"
	"sigs.k8s.io/kustomize/kyaml/yaml"
)

var (
	deploymentResourceManifest = `
apiVersion: apps/v1
kind: Deployment
metadata:
  namespace: myspace
  name: mysql-deployment
spec:
  replicas: 3
  foo: bar
  template:
    spec:
      containers:
      - name: mysql
        image: mysql:1.7.9
`

	configMapResourceManifest = `
apiVersion: v1
kind: ConfigMap
metadata:
  name: configmap
data:
  foo: bar
`
)

var (
	DeploymentResource = "deployment"
	ConfigMapResource  = "configmap"
	resources          = map[string]resourceInfo{
		DeploymentResource: {
			filename: "deployment.yaml",
			manifest: deploymentResourceManifest,
		},
		ConfigMapResource: {
			filename: "configmap.yaml",
			manifest: configMapResourceManifest,
		},
	}
)

type resourceInfo struct {
	filename string
	manifest string
}

type resourceInfoWithSetters struct {
	resourceInfo resourceInfo
	setterRefs   []SetterRef
	mutators     []yaml.Filter
}

// Pkg represents a package that can be created on the file system
// by using the Build function
type Pkg struct {
	Name string

	Kptfile *Kptfile

	resources []resourceInfoWithSetters

	files map[string]string

	subPkgs []*Pkg
}

func NewKptfile() *Kptfile {
	return &Kptfile{}
}

// Kptfile represents the Kptfile of a package.
type Kptfile struct {
	Setters []Setter
	Repo    string
	Ref     string
}

// WithUpstream adds information about the upstream information to the Kptfile.
// The upstream section of the Kptfile is only added if this information is
// provided.
func (k *Kptfile) WithUpstream(repo, ref string) *Kptfile {
	k.Repo = repo
	k.Ref = ref
	return k
}

// WithSetters adds information about the setters for a Kptfile.
func (k *Kptfile) WithSetters(setters ...Setter) *Kptfile {
	k.Setters = setters
	return k
}

// Setter contains the properties required for adding a setter to the
// Kptfile.
type Setter struct {
	Name  string
	Value string
	IsSet bool
}

// NewSetter creates a new setter that is not marked as set
func NewSetter(name, value string) Setter {
	return Setter{
		Name:  name,
		Value: value,
	}
}

// NewSetSetter creates a new setter that is marked as set.
func NewSetSetter(name, value string) Setter {
	return Setter{
		Name:  name,
		Value: value,
		IsSet: true,
	}
}

// SetterRef specifies the information for creating a new reference to
// a setter in a resource.
type SetterRef struct {
	Path []string
	Name string
}

// NewSetterRef creates a new setterRef with the given name and path.
func NewSetterRef(name string, path ...string) SetterRef {
	return SetterRef{
		Path: path,
		Name: name,
	}
}

// NewPackage creates a new package for testing.
func NewPackage(name string) *Pkg {
	return &Pkg{
		Name:  name,
		files: make(map[string]string),
	}
}

// WithKptfile configures the current package to have a Kptfile. Only
// zero or one Kptfiles are accepted.
func (p *Pkg) WithKptfile(kf ...*Kptfile) *Pkg {
	if len(kf) > 1 {
		panic("only 0 or 1 Kptfiles are allowed")
	}
	if len(kf) == 0 {
		p.Kptfile = NewKptfile()
	} else {
		p.Kptfile = kf[0]
	}
	return p
}

// WithResource configures the package to include the provided resource
func (p *Pkg) WithResource(resourceName string, mutators ...yaml.Filter) *Pkg {
	resourceInfo, ok := resources[resourceName]
	if !ok {
		panic(fmt.Errorf("unknown resource %s", resourceName))
	}
	p.resources = append(p.resources, resourceInfoWithSetters{
		resourceInfo: resourceInfo,
		setterRefs:   []SetterRef{},
		mutators:     mutators,
	})
	return p
}

// WithResourceAndSetters configures the package to have the provided resource.
// It also allows for specifying setterRefs for the resource and a set of
// mutators that will update the content of the resource.
func (p *Pkg) WithResourceAndSetters(resourceName string, setterRefs []SetterRef, mutators ...yaml.Filter) *Pkg {
	resourceInfo, ok := resources[resourceName]
	if !ok {
		panic(fmt.Errorf("unknown resource %s", resourceName))
	}
	p.resources = append(p.resources, resourceInfoWithSetters{
		resourceInfo: resourceInfo,
		setterRefs:   setterRefs,
		mutators:     mutators,
	})
	return p
}

func (p *Pkg) WithFile(name, content string) *Pkg {
	p.files[name] = content
	return p
}

// WithSubPackages adds the provided packages as subpackages to the current
// package
func (p *Pkg) WithSubPackages(ps ...*Pkg) *Pkg {
	p.subPkgs = append(p.subPkgs, ps...)
	return p
}

// Build outputs the current data structure as a set of (nested) package
// in the provided path.
func (p *Pkg) Build(path string) error {
	_, err := os.Stat(path)
	if err != nil {
		return err
	}
	return buildRecursive(path, p)
}

func buildRecursive(path string, pkg *Pkg) error {
	pkgPath := filepath.Join(path, pkg.Name)
	err := os.Mkdir(pkgPath, 0700)
	if err != nil {
		return err
	}

	if pkg.Kptfile != nil {
		content := buildKptfile(pkg)

		err := ioutil.WriteFile(filepath.Join(pkgPath, kptfileutil.KptFileName),
			[]byte(content), 0600)
		if err != nil {
			return err
		}
	}

	for _, ri := range pkg.resources {
		m := ri.resourceInfo.manifest
		r := yaml.MustParse(m)
		for _, setterRef := range ri.setterRefs {
			n, err := r.Pipe(yaml.PathGetter{
				Path: setterRef.Path,
			})
			if err != nil {
				return err
			}
			n.YNode().LineComment = fmt.Sprintf(`{"$openapi":"%s"}`, setterRef.Name)
		}

		for _, m := range ri.mutators {
			if err := r.PipeE(m); err != nil {
				return err
			}
		}

		filePath := filepath.Join(pkgPath, ri.resourceInfo.filename)
		err = ioutil.WriteFile(filePath, []byte(r.MustString()), 0600)
		if err != nil {
			return err
		}
	}

	for name, content := range pkg.files {
		filePath := filepath.Join(pkgPath, name)
		_, err := os.Stat(filePath)
		if err != nil && !os.IsNotExist(err) {
			return err
		}
		if !os.IsNotExist(err) {
			return fmt.Errorf("file %s already exists", name)
		}
		err = ioutil.WriteFile(filePath, []byte(content), 0600)
		if err != nil {
			return err
		}
	}

	for i := range pkg.subPkgs {
		subPkg := pkg.subPkgs[i]
		err = buildRecursive(pkgPath, subPkg)
		if err != nil {
			return err
		}
	}

	return nil
}

var kptfileTemplate = `
apiVersion: kpt.dev/v1alpha1
kind: Kptfile
metadata:
  name: {{.Name}}
{{- if gt (len .Kptfile.Setters) 0 }}
openAPI:
  definitions:
{{- range .Kptfile.Setters }}
    io.k8s.cli.setters.{{.Name}}:
      x-k8s-cli:
        setter:
          name: {{.Name}}
          value: {{.Value}}
{{- if eq .IsSet true }}
          isSet: true
{{- end }}
{{- end }}
{{- end }}
{{- if gt (len .Kptfile.Repo) 0 }}
upstream:
  type: git
  git:
    ref: {{.Kptfile.Ref}}
    repo: {{.Kptfile.Repo}}
{{- end }}
`

func buildKptfile(pkg *Pkg) string {
	tmpl, err := template.New("test").Parse(kptfileTemplate)
	if err != nil {
		panic(err)
	}
	var buf bytes.Buffer
	err = tmpl.Execute(&buf, pkg)
	if err != nil {
		panic(err)
	}
	result := buf.String()
	return result
}

func ExpandPkg(t *testing.T, pkg *Pkg) string {
	if pkg.Name == "" {
		pkg.Name = "base"
	}
	dir, err := ioutil.TempDir("", "test-kpt-builder-")
	if !assert.NoError(t, err) {
		t.FailNow()
	}
	err = pkg.Build(dir)
	if !assert.NoError(t, err) {
		t.FailNow()
	}
	return filepath.Join(dir, pkg.Name)
}
<|endoftext|>


<commit_before>package pqtsql

import (
	"bytes"
	"errors"
	"fmt"
	"io"

	"github.com/piotrkowalczuk/pqt"
)

// Generator ...
type Generator struct {
	Version float64
}

// Generate ...
func (g *Generator) Generate(s *pqt.Schema) ([]byte, error) {
	code, err := g.generate(s)
	if err != nil {
		return nil, err
	}

	return code.Bytes(), nil
}

// GenerateTo ...
func (g *Generator) GenerateTo(s *pqt.Schema, w io.Writer) error {
	code, err := g.generate(s)
	if err != nil {
		return err
	}

	_, err = code.WriteTo(w)
	return err
}

func (g *Generator) generate(s *pqt.Schema) (*bytes.Buffer, error) {
	code := bytes.NewBufferString("-- do not modify, generated by pqt\n\n")
	if s.Name != "" {
		fmt.Fprint(code, "CREATE SCHEMA ")
		if s.IfNotExists {
			fmt.Fprint(code, "IF NOT EXISTS ")
		}
		fmt.Fprintf(code, "%s; \n\n", s.Name)
	}
	for _, f := range s.Functions {
		if err := g.generateCreateFunction(code, f); err != nil {
			return nil, err
		}
	}
	for _, t := range s.Tables {
		if err := g.generateCreateTable(code, t); err != nil {
			return nil, err
		}
		for _, cnstr := range t.Constraints {
			switch cnstr.Type {
			case pqt.ConstraintTypeIndex:
				indexConstraintQuery(code, cnstr, g.Version)
			case pqt.ConstraintTypeUniqueIndex:
				uniqueIndexConstraintQuery(code, cnstr, g.Version)
			}
		}
		fmt.Fprintln(code, "")
	}

	return code, nil
}

func (g *Generator) generateCreateFunction(buf *bytes.Buffer, f *pqt.Function) error {
	if f == nil {
		return nil
	}
	if f.Name == "" {
		return errors.New("missing function name")
	}

	buf.WriteString("CREATE OR REPLACE FUNCTION ")
	buf.WriteString(f.Name)
	buf.WriteString("(")
	for i, arg := range f.Args {
		if i != 0 {
			buf.WriteString(", ")
		}
		buf.WriteString(arg.Name)
		buf.WriteString(" ")
		buf.WriteString(arg.Type.String())
	}
	buf.WriteString(") RETURNS ")
	buf.WriteString(f.Type.String())
	buf.WriteString("\n	AS '")
	buf.WriteString(f.Body)
	buf.WriteString("'\n	LANGUAGE SQL")
	switch f.Behaviour {
	case pqt.FunctionBehaviourVolatile:
		buf.WriteString("\n	VOLATILE")
	case pqt.FunctionBehaviourImmutable:
		buf.WriteString("\n	IMMUTABLE")
	case pqt.FunctionBehaviourStable:
		buf.WriteString("\n	STABLE")
	}
	buf.WriteString(";\n\n")

	return nil
}

func (g *Generator) generateCreateTable(buf *bytes.Buffer, t *pqt.Table) error {
	if t == nil {
		return nil
	}

	if t.Name == "" {
		return errors.New("missing table name")
	}
	if len(t.Columns) == 0 {
		return fmt.Errorf("table %s has no columns", t.Name)
	}

	buf.WriteString("CREATE ")
	if t.Temporary {
		buf.WriteString("TEMPORARY ")
	}
	buf.WriteString("TABLE ")
	if t.IfNotExists {
		buf.WriteString("IF NOT EXISTS ")
	}
	if t.Schema != nil {
		buf.WriteString(t.Schema.Name)
		buf.WriteRune('.')
		buf.WriteString(t.Name)
	} else {
		buf.WriteString(t.Name)
	}
	buf.WriteString(" (\n")

	constraints := t.Constraints
	for _, r := range t.OwnedRelationships {
		// If ...
		if len(r.OwnerColumns) == 1 {
			continue
		}
		if r.OwnerForeignKey != nil {
			constraints = append(constraints, r.OwnerForeignKey)
		}
	}
	nbOfConstraints := constraints.CountOf(
		pqt.ConstraintTypePrimaryKey,
		pqt.ConstraintTypeCheck,
		pqt.ConstraintTypeUnique,
		pqt.ConstraintTypeForeignKey,
		pqt.ConstraintTypeExclusion,
	)
	for i, c := range t.Columns {
		if c.IsDynamic {
			continue
		}
		buf.WriteRune('	')
		buf.WriteString(c.Name)
		buf.WriteRune(' ')
		buf.WriteString(c.Type.String())
		if c.Collate != "" {
			buf.WriteRune(' ')
			buf.WriteString(c.Collate)
		}
		if d, ok := c.DefaultOn(pqt.EventInsert); ok {
			buf.WriteString(" DEFAULT ")
			buf.WriteString(d)
		}
		if c.NotNull {
			buf.WriteString(" NOT NULL")
		}

		if i < len(t.Columns)-1 || nbOfConstraints > 0 {
			buf.WriteRune(',')
		}
		buf.WriteRune('\n')
	}

	if nbOfConstraints > 0 {
		buf.WriteRune('\n')
	}

	i := 0
	for _, c := range constraints {
		if c.Type == pqt.ConstraintTypeIndex || c.Type == pqt.ConstraintTypeUniqueIndex {
			continue
		}
		buf.WriteString("	")
		err := g.generateConstraint(buf, c)
		if err != nil {
			return err
		}
		if i < nbOfConstraints-1 {
			buf.WriteRune(',')
		}
		buf.WriteRune('\n')
		i++
	}

	buf.WriteString(");\n")

	return nil
}

func (g *Generator) generateConstraint(buf *bytes.Buffer, c *pqt.Constraint) error {
	switch c.Type {
	case pqt.ConstraintTypeUnique:
		uniqueConstraintQuery(buf, c)
	case pqt.ConstraintTypePrimaryKey:
		primaryKeyConstraintQuery(buf, c)
	case pqt.ConstraintTypeForeignKey:
		return foreignKeyConstraintQuery(buf, c)
	case pqt.ConstraintTypeCheck:
		checkConstraintQuery(buf, c)
	case pqt.ConstraintTypeIndex:
	case pqt.ConstraintTypeUniqueIndex:
	default:
		return fmt.Errorf("unknown constraint type: %s", c.Type)
	}

	return nil
}

func uniqueConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint) {
	fmt.Fprintf(buf, `CONSTRAINT "%s" UNIQUE (%s)`, c.Name(), pqt.JoinColumns(c.PrimaryColumns, ", "))
}

func primaryKeyConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint) {
	fmt.Fprintf(buf, `CONSTRAINT "%s" PRIMARY KEY (%s)`, c.Name(), pqt.JoinColumns(c.PrimaryColumns, ", "))
}

func foreignKeyConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint) error {
	switch {
	case len(c.PrimaryColumns) == 0:
		return errors.New("foreign key constraint require at least one column")
	case len(c.Columns) == 0:
		return errors.New("foreign key constraint require at least one reference column")
	case c.Table == nil:
		return errors.New("foreiqn key constraint missing reference table")
	}

	fmt.Fprintf(buf, `CONSTRAINT "%s" FOREIGN KEY (%s) REFERENCES %s (%s)`,
		c.Name(),
		pqt.JoinColumns(c.PrimaryColumns, ", "),
		c.Table.FullName(),
		pqt.JoinColumns(c.Columns, ", "),
	)

	switch c.OnDelete {
	case pqt.Cascade:
		buf.WriteString(" ON DELETE CASCADE")
	case pqt.Restrict:
		buf.WriteString(" ON DELETE RESTRICT")
	case pqt.SetNull:
		buf.WriteString(" ON DELETE SET NULL")
	case pqt.SetDefault:
		buf.WriteString(" ON DELETE SET DEFAULT")
	}

	switch c.OnUpdate {
	case pqt.Cascade:
		buf.WriteString(" ON UPDATE CASCADE")
	case pqt.Restrict:
		buf.WriteString(" ON UPDATE RESTRICT")
	case pqt.SetNull:
		buf.WriteString(" ON UPDATE SET NULL")
	case pqt.SetDefault:
		buf.WriteString(" ON UPDATE SET DEFAULT")
	}

	return nil
}

func checkConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint) {
	fmt.Fprintf(buf, `CONSTRAINT "%s" CHECK (%s)`, c.Name(), c.Check)
}

func indexConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint, ver float64) {
	// TODO: change code so IF NOT EXISTS is optional
	if ver >= 9.5 {
		fmt.Fprintf(buf, `CREATE INDEX IF NOT EXISTS "%s" ON %s (%s);`, c.Name(), c.PrimaryTable.FullName(), c.PrimaryColumns.String())
	} else {
		fmt.Fprintf(buf, `CREATE INDEX "%s" ON %s (%s);`, c.Name(), c.PrimaryTable.FullName(), c.PrimaryColumns.String())
	}
	fmt.Fprintln(buf, "")
}

func uniqueIndexConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint, ver float64) {
	fmt.Fprintf(buf, `CREATE UNIQUE INDEX "%s" ON %s (%s)`, c.Name(), c.PrimaryTable.FullName(), c.PrimaryColumns.String())

	if c.Where != "" {
		fmt.Fprintf(buf, " WHERE %s", c.Where)
	}

	fmt.Fprint(buf, ";\n")
}
<commit_msg>'if not exists' on unique index<commit_after>package pqtsql

import (
	"bytes"
	"errors"
	"fmt"
	"io"

	"github.com/piotrkowalczuk/pqt"
)

// Generator ...
type Generator struct {
	Version float64
}

// Generate ...
func (g *Generator) Generate(s *pqt.Schema) ([]byte, error) {
	code, err := g.generate(s)
	if err != nil {
		return nil, err
	}

	return code.Bytes(), nil
}

// GenerateTo ...
func (g *Generator) GenerateTo(s *pqt.Schema, w io.Writer) error {
	code, err := g.generate(s)
	if err != nil {
		return err
	}

	_, err = code.WriteTo(w)
	return err
}

func (g *Generator) generate(s *pqt.Schema) (*bytes.Buffer, error) {
	code := bytes.NewBufferString("-- do not modify, generated by pqt\n\n")
	if s.Name != "" {
		fmt.Fprint(code, "CREATE SCHEMA ")
		if s.IfNotExists {
			fmt.Fprint(code, "IF NOT EXISTS ")
		}
		fmt.Fprintf(code, "%s; \n\n", s.Name)
	}
	for _, f := range s.Functions {
		if err := g.generateCreateFunction(code, f); err != nil {
			return nil, err
		}
	}
	for _, t := range s.Tables {
		if err := g.generateCreateTable(code, t); err != nil {
			return nil, err
		}
		for _, cnstr := range t.Constraints {
			switch cnstr.Type {
			case pqt.ConstraintTypeIndex:
				indexConstraintQuery(code, cnstr, g.Version)
			case pqt.ConstraintTypeUniqueIndex:
				uniqueIndexConstraintQuery(code, cnstr, g.Version)
			}
		}
		fmt.Fprintln(code, "")
	}

	return code, nil
}

func (g *Generator) generateCreateFunction(buf *bytes.Buffer, f *pqt.Function) error {
	if f == nil {
		return nil
	}
	if f.Name == "" {
		return errors.New("missing function name")
	}

	buf.WriteString("CREATE OR REPLACE FUNCTION ")
	buf.WriteString(f.Name)
	buf.WriteString("(")
	for i, arg := range f.Args {
		if i != 0 {
			buf.WriteString(", ")
		}
		buf.WriteString(arg.Name)
		buf.WriteString(" ")
		buf.WriteString(arg.Type.String())
	}
	buf.WriteString(") RETURNS ")
	buf.WriteString(f.Type.String())
	buf.WriteString("\n	AS '")
	buf.WriteString(f.Body)
	buf.WriteString("'\n	LANGUAGE SQL")
	switch f.Behaviour {
	case pqt.FunctionBehaviourVolatile:
		buf.WriteString("\n	VOLATILE")
	case pqt.FunctionBehaviourImmutable:
		buf.WriteString("\n	IMMUTABLE")
	case pqt.FunctionBehaviourStable:
		buf.WriteString("\n	STABLE")
	}
	buf.WriteString(";\n\n")

	return nil
}

func (g *Generator) generateCreateTable(buf *bytes.Buffer, t *pqt.Table) error {
	if t == nil {
		return nil
	}

	if t.Name == "" {
		return errors.New("missing table name")
	}
	if len(t.Columns) == 0 {
		return fmt.Errorf("table %s has no columns", t.Name)
	}

	buf.WriteString("CREATE ")
	if t.Temporary {
		buf.WriteString("TEMPORARY ")
	}
	buf.WriteString("TABLE ")
	if t.IfNotExists {
		buf.WriteString("IF NOT EXISTS ")
	}
	if t.Schema != nil {
		buf.WriteString(t.Schema.Name)
		buf.WriteRune('.')
		buf.WriteString(t.Name)
	} else {
		buf.WriteString(t.Name)
	}
	buf.WriteString(" (\n")

	constraints := t.Constraints
	for _, r := range t.OwnedRelationships {
		// If ...
		if len(r.OwnerColumns) == 1 {
			continue
		}
		if r.OwnerForeignKey != nil {
			constraints = append(constraints, r.OwnerForeignKey)
		}
	}
	nbOfConstraints := constraints.CountOf(
		pqt.ConstraintTypePrimaryKey,
		pqt.ConstraintTypeCheck,
		pqt.ConstraintTypeUnique,
		pqt.ConstraintTypeForeignKey,
		pqt.ConstraintTypeExclusion,
	)
	for i, c := range t.Columns {
		if c.IsDynamic {
			continue
		}
		buf.WriteRune('	')
		buf.WriteString(c.Name)
		buf.WriteRune(' ')
		buf.WriteString(c.Type.String())
		if c.Collate != "" {
			buf.WriteRune(' ')
			buf.WriteString(c.Collate)
		}
		if d, ok := c.DefaultOn(pqt.EventInsert); ok {
			buf.WriteString(" DEFAULT ")
			buf.WriteString(d)
		}
		if c.NotNull {
			buf.WriteString(" NOT NULL")
		}

		if i < len(t.Columns)-1 || nbOfConstraints > 0 {
			buf.WriteRune(',')
		}
		buf.WriteRune('\n')
	}

	if nbOfConstraints > 0 {
		buf.WriteRune('\n')
	}

	i := 0
	for _, c := range constraints {
		if c.Type == pqt.ConstraintTypeIndex || c.Type == pqt.ConstraintTypeUniqueIndex {
			continue
		}
		buf.WriteString("	")
		err := g.generateConstraint(buf, c)
		if err != nil {
			return err
		}
		if i < nbOfConstraints-1 {
			buf.WriteRune(',')
		}
		buf.WriteRune('\n')
		i++
	}

	buf.WriteString(");\n")

	return nil
}

func (g *Generator) generateConstraint(buf *bytes.Buffer, c *pqt.Constraint) error {
	switch c.Type {
	case pqt.ConstraintTypeUnique:
		uniqueConstraintQuery(buf, c)
	case pqt.ConstraintTypePrimaryKey:
		primaryKeyConstraintQuery(buf, c)
	case pqt.ConstraintTypeForeignKey:
		return foreignKeyConstraintQuery(buf, c)
	case pqt.ConstraintTypeCheck:
		checkConstraintQuery(buf, c)
	case pqt.ConstraintTypeIndex:
	case pqt.ConstraintTypeUniqueIndex:
	default:
		return fmt.Errorf("unknown constraint type: %s", c.Type)
	}

	return nil
}

func uniqueConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint) {
	fmt.Fprintf(buf, `CONSTRAINT "%s" UNIQUE (%s)`, c.Name(), pqt.JoinColumns(c.PrimaryColumns, ", "))
}

func primaryKeyConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint) {
	fmt.Fprintf(buf, `CONSTRAINT "%s" PRIMARY KEY (%s)`, c.Name(), pqt.JoinColumns(c.PrimaryColumns, ", "))
}

func foreignKeyConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint) error {
	switch {
	case len(c.PrimaryColumns) == 0:
		return errors.New("foreign key constraint require at least one column")
	case len(c.Columns) == 0:
		return errors.New("foreign key constraint require at least one reference column")
	case c.Table == nil:
		return errors.New("foreiqn key constraint missing reference table")
	}

	fmt.Fprintf(buf, `CONSTRAINT "%s" FOREIGN KEY (%s) REFERENCES %s (%s)`,
		c.Name(),
		pqt.JoinColumns(c.PrimaryColumns, ", "),
		c.Table.FullName(),
		pqt.JoinColumns(c.Columns, ", "),
	)

	switch c.OnDelete {
	case pqt.Cascade:
		buf.WriteString(" ON DELETE CASCADE")
	case pqt.Restrict:
		buf.WriteString(" ON DELETE RESTRICT")
	case pqt.SetNull:
		buf.WriteString(" ON DELETE SET NULL")
	case pqt.SetDefault:
		buf.WriteString(" ON DELETE SET DEFAULT")
	}

	switch c.OnUpdate {
	case pqt.Cascade:
		buf.WriteString(" ON UPDATE CASCADE")
	case pqt.Restrict:
		buf.WriteString(" ON UPDATE RESTRICT")
	case pqt.SetNull:
		buf.WriteString(" ON UPDATE SET NULL")
	case pqt.SetDefault:
		buf.WriteString(" ON UPDATE SET DEFAULT")
	}

	return nil
}

func checkConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint) {
	fmt.Fprintf(buf, `CONSTRAINT "%s" CHECK (%s)`, c.Name(), c.Check)
}

func indexConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint, ver float64) {
	// TODO: change code so IF NOT EXISTS is optional
	if ver >= 9.5 {
		fmt.Fprintf(buf, `CREATE INDEX IF NOT EXISTS "%s" ON %s (%s);`, c.Name(), c.PrimaryTable.FullName(), c.PrimaryColumns.String())
	} else {
		fmt.Fprintf(buf, `CREATE INDEX "%s" ON %s (%s);`, c.Name(), c.PrimaryTable.FullName(), c.PrimaryColumns.String())
	}
	fmt.Fprintln(buf, "")
}

func uniqueIndexConstraintQuery(buf *bytes.Buffer, c *pqt.Constraint, ver float64) {
	if ver >= 9.5 {
		fmt.Fprintf(buf, `CREATE UNIQUE INDEX IF NOT EXISTS "%s" ON %s (%s)`, c.Name(), c.PrimaryTable.FullName(), c.PrimaryColumns.String())
	} else {
		fmt.Fprintf(buf, `CREATE UNIQUE INDEX "%s" ON %s (%s)`, c.Name(), c.PrimaryTable.FullName(), c.PrimaryColumns.String())
	}
	if c.Where != "" {
		fmt.Fprintf(buf, " WHERE %s", c.Where)
	}

	fmt.Fprint(buf, ";\n")
}
<|endoftext|>


<commit_before>/**
 * ws.go - A light weight webserver for static content
 * development and prototyping route based web API.
 *
 * Supports both http and https protocols. Dynamic route
 * processing available via Otto JavaScript virtual machines.
 *
 * @author R. S. Doiel, <rsdoiel@yahoo.com>
 * copyright (c) 2014
 * All rights reserved.
 * @license BSD 2-Clause License
 */
package main

import (
	"./ottoengine"
	"./wslog"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	//"errors"
	"flag"
	"fmt"
	"log"
	"math/big"
	"net"
	"net/http"
	"os"
	"os/user"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

var REVISION = "v0.0.0-alpha"

// variables for keygen
var (
	cli_keygen   = flag.Bool("keygen", false, "Generate TLS ceriticates and keys")
	cli_ssl_host = flag.String("keygen-ssl-host", "", "Comma-separated hostnames and IPs to generate a certificate for")
	validFrom    = flag.String("keygen-start-date", "", "Creation date formatted as Jan 1 15:04:05 2011")
	validFor     = flag.Duration("keygen-duration", 365*24*time.Hour, "Duration that certificate is valid for")
	organization = flag.String("keygen-organization", "Acme Co.", "Organization used to sign certificate")
	isCA         = flag.Bool("keygen-ca", false, "whether this cert should be its own Certificate Authority")
	rsaBits      = flag.Int("keygen-rsa-bits", 2048, "Size of RSA key to generate")
)

// command line parameters that override environment variables
var (
	cli_use_tls   *bool
	cli_docroot   *string
	cli_host      *string
	cli_port      *string
	cli_cert      *string
	cli_key       *string
	cli_otto      *bool
	cli_otto_path *string
	cli_version   *bool
)

var Usage = func() {
	flag.PrintDefaults()
}

// Application's profile - who started the process, port assignment
// configuration settings, etc.
type Profile struct {
	Username  string
	Hostname  string
	Port      string
	Use_TLS   bool
	Docroot   string
	Cert      string
	Key       string
	Otto      bool
	Otto_Path string
}

func LoadProfile(cli_docroot string, cli_host string, cli_port string, cli_use_tls bool, cli_cert string, cli_key string, cli_otto bool, cli_otto_path string) (*Profile, error) {
	ws_user, err := user.Current()
	if err != nil {
		return nil, err
	}
	hostname, err := os.Hostname()
	if err != nil {
		return nil, err
	}
	port := "8000"
	use_tls := false
	otto := false
	otto_path := ""

	cert := ""
	key := ""
	if err != nil {
		return nil, err
	}
	docroot, _ := os.Getwd()

	// now overwrite with any environment settings found.
	env_host := os.Getenv("WS_HOST")
	env_port := os.Getenv("WS_PORT")
	env_use_tls := os.Getenv("WS_TLS")
	env_cert := os.Getenv("WS_CERT")
	env_key := os.Getenv("WS_KEY")
	env_docroot := os.Getenv("WS_DOCROOT")
	env_otto := os.Getenv("WS_OTTO")
	env_otto_path := os.Getenv("WS_OTTO_PATH")
	if env_host != "" {
		hostname = env_host
	}
	if env_use_tls == "true" {
		use_tls = true
		port = "8443"
	}
	if env_port != "" {
		port = env_port
	}
	if env_docroot != "" {
		docroot = env_docroot
	}
	if env_cert != "" {
		cert = env_cert
	}
	if env_key != "" {
		key = env_key
	}
	if env_otto == "true" {
		otto = true
	}
	if env_otto_path != "" {
		otto_path = env_otto_path
	}

	// Finally resolve any command line overrides
	if cli_docroot != "" {
		docroot = cli_docroot
	}
	if cli_use_tls == true {
		use_tls = true
		if env_port == "" {
			port = "8443"
		}
	}
	if cli_host != "" {
		hostname = cli_host
	}
	if cli_port != "" {
		port = cli_port
	}
	if cli_cert != "" {
		cert = cli_cert
	}
	if cli_key != "" {
		key = cli_key
	}
	if cli_otto == true {
		otto = true
	}
	if cli_otto_path != "" {
		otto_path = cli_otto_path
	}

	// If TLS is false then don't expose the location of the cert/key
	if use_tls == false {
		cert = ""
		key = ""
	}

	// Normalize docroot
	if strings.HasPrefix(docroot, "/") == false {
		clean_docroot, err := filepath.Abs(path.Join("./", docroot))
		if err != nil {
			log.Fatalf("Can't expand docroot %s: %s\n", docroot, err)
		}
		docroot = clean_docroot
	}
	// Normalize otto_path
	if strings.HasPrefix(otto_path, "/") == false {
		clean_otto_path, err := filepath.Abs(path.Join("./", otto_path))
		if err != nil {
			log.Fatalf("Can't expand otto_path %s: %s\n", otto_path, err)
		}
		otto_path = clean_otto_path
	}
	return &Profile{
		Username:  ws_user.Username,
		Hostname:  hostname,
		Port:      port,
		Docroot:   docroot,
		Use_TLS:   use_tls,
		Cert:      cert,
		Key:       key,
		Otto:      otto,
		Otto_Path: otto_path}, nil
}

func request_log(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		wslog.LogRequest(r.Method, r.URL, r.RemoteAddr, r.Proto, r.Referer(), r.UserAgent())
		handler.ServeHTTP(w, r)
	})
}

func Webserver(profile *Profile) error {
	// If otto is enabled add routes and handle them.
	if profile.Otto == true {
		otto_path, err := filepath.Abs(profile.Otto_Path)
		if err != nil {
			log.Fatalf("Can't read %s: %s\n", profile.Otto_Path, err)
		}
		programs, err := ottoengine.Load(otto_path)
		if err != nil {
			log.Fatalf("Load error: %s\n", err)
		}
		ottoengine.AddRoutes(programs)
	}

	// Restricted FileService excluding dot files and directories
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		var hasDotPath = regexp.MustCompile(`\/\.`)
		unclean_path := r.URL.Path
		if !strings.HasPrefix(unclean_path, "/") {
			unclean_path = "/" + unclean_path
		}
		clean_path := path.Clean(unclean_path)
		r.URL.Path = clean_path
		resolved_path := path.Clean(path.Join(profile.Docroot, clean_path))
		_, err := os.Stat(resolved_path)
		if hasDotPath.MatchString(clean_path) == true ||
			strings.HasPrefix(resolved_path, profile.Docroot) == false ||
			os.IsPermission(err) == true {
			wslog.LogResponse(401, "Not Authorized", r.Method, r.URL, r.RemoteAddr, resolved_path, "")
			http.Error(w, "Not Authorized", 401)
		} else if os.IsNotExist(err) == true {
			wslog.LogResponse(404, "Not Found", r.Method, r.URL, r.RemoteAddr, resolved_path, "")
			http.NotFound(w, r)
		} else if err == nil {
			wslog.LogResponse(200, "OK", r.Method, r.URL, r.RemoteAddr, resolved_path, "")
			http.ServeFile(w, r, resolved_path)
		} else {
			// Easter egg
			wslog.LogResponse(418, "I'm a teapot", r.Method, r.URL, r.RemoteAddr, resolved_path, "")
			http.Error(w, "I'm a teapot", 418)
		}
	})

	// Now start up the server and log transactions
	if profile.Use_TLS == true {
		if profile.Cert == "" || profile.Key == "" {
			log.Fatalf("TLS set true but missing key or certificate")
		}
		log.Println("Starting https://" + net.JoinHostPort(profile.Hostname, profile.Port))
		return http.ListenAndServeTLS(net.JoinHostPort(profile.Hostname, profile.Port), profile.Cert, profile.Key, request_log(http.DefaultServeMux))
	}
	log.Println("Starting http://" + net.JoinHostPort(profile.Hostname, profile.Port))
	// Now start up the server and log transactions
	return http.ListenAndServe(net.JoinHostPort(profile.Hostname, profile.Port), request_log(http.DefaultServeMux))
}

func keygen(profile *Profile) error {
	home := os.Getenv("HOME")
	certFilename := profile.Cert
	if certFilename == "" {
		certFilename = path.Join(home, "etc/ws/cert.pem")
	}
	keyFilename := profile.Key
	if keyFilename == "" {
		keyFilename = path.Join(home, "etc/ws/key.pem")
	}

	hostnames := profile.Hostname
	if *cli_ssl_host != "" {
		hostnames = *cli_ssl_host
	}
	if hostnames == "" {
		log.Fatalf("Missing required -ssl-host parameter")
	}

	log.Printf("\n\n"+
		" Cert: %s\n"+
		"  Key: %s\n"+
		" Host: %s\n"+
		" Organization: %v\n"+
		"\n\n",
		certFilename,
		keyFilename,
		hostnames,
		*organization)

	priv, err := rsa.GenerateKey(rand.Reader, *rsaBits)
	if err != nil {
		log.Fatalf("failed to generate private key: %s", err)
	}

	var notBefore time.Time
	if len(*validFrom) == 0 {
		notBefore = time.Now()
	} else {
		notBefore, err = time.Parse("Jan 2 15:04:05 2006", *validFrom)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to parse creation date: %s\n", err)
			return err
		}
	}

	notAfter := notBefore.Add(*validFor)

	// end of ASN.1 time
	endOfTime := time.Date(2049, 12, 31, 2, 59, 59, 0, time.UTC)
	if notAfter.After(endOfTime) {
		notAfter = endOfTime
	}

	template := x509.Certificate{
		SerialNumber: new(big.Int).SetInt64(0),
		Subject: pkix.Name{
			Organization: []string{*organization},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	hosts := strings.Split(hostnames, ",")
	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, h)
		}
	}

	if *isCA {
		template.IsCA = true
		template.KeyUsage |= x509.KeyUsageCertSign
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		log.Fatalf("Failed to create certificate: %s", err)
	}

	certOut, err := os.Create(certFilename)
	if err != nil {
		log.Fatalf("failed to open cert.pem for writing: %s", err)
	}
	pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	certOut.Close()
	log.Printf("written %s\n", certFilename)

	keyOut, err := os.OpenFile(keyFilename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		log.Print("failed to open key.pem for writing:", err)
		return err
	}
	pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})
	keyOut.Close()
	log.Printf("written %s\n", keyFilename)
	// We got this for so no errors
	return nil
}

func init() {
	// command line parameters that override environment variables, many have short forms too.
	shortform := " (short form)"

	// No short form
	cli_use_tls = flag.Bool("tls", false, "Turn on TLS (https) support with true, off with false (default is false)")
	cli_key = flag.String("key", "", "path to your SSL key pem file.")
	cli_cert = flag.String("cert", "", "path to your SSL cert pem file.")

	// These have short forms too
	msg := "document root"
	cli_docroot = flag.String("docroot", "", msg)
	flag.StringVar(cli_docroot, "D", "", msg+shortform)

	msg = "hostname for webserver"
	cli_host = flag.String("host", "", msg)
	flag.StringVar(cli_host, "H", "", msg+shortform)

	msg = "Port number to listen on"
	cli_port = flag.String("port", "", msg)
	flag.StringVar(cli_port, "P", "", msg+shortform)

	msg = "turn on ottoengine, defaults to false"
	cli_otto = flag.Bool("otto", false, msg)
	flag.BoolVar(cli_otto, "o", false, msg+shortform)

	msg = "directory containingo your ottoengine JavaScript files"
	cli_otto_path = flag.String("otto-path", "", msg)
	flag.StringVar(cli_otto_path, "O", "", msg+shortform)

	msg = "Display the version number"
	cli_version = flag.Bool("version", false, msg)
	flag.BoolVar(cli_version, "v", false, msg+shortform)
}

func main() {
	flag.Parse()

	if *cli_version == true {
		fmt.Println(REVISION)
		os.Exit(0)
	}

	profile, _ := LoadProfile(*cli_docroot, *cli_host, *cli_port, *cli_use_tls, *cli_cert, *cli_key, *cli_otto, *cli_otto_path)
	if *cli_keygen == true {
		err := keygen(profile)
		if err != nil {
			log.Fatalf("%s\n", err)
		}
		os.Exit(0)
	}

	log.Printf("\n\n"+
		"          TLS: %t\n"+
		"         Cert: %s\n"+
		"          Key: %s\n"+
		"      Docroot: %s\n"+
		"         Host: %s\n"+
		"         Port: %s\n"+
		"       Run as: %s\n\n"+
		" Otto enabled: %t\n"+
		"         Path: %s\n"+
		"\n\n",
		profile.Use_TLS,
		profile.Cert,
		profile.Key,
		profile.Docroot,
		profile.Hostname,
		profile.Port,
		profile.Username,
		profile.Otto,
		profile.Otto_Path)
	err := Webserver(profile)
	if err != nil {
		log.Fatal(err)
	}
	os.Exit(0)
}
<commit_msg>orginizing code<commit_after>/**
 * ws.go - A light weight webserver for static content
 * development and prototyping route based web API.
 *
 * Supports both http and https protocols. Dynamic route
 * processing available via Otto JavaScript virtual machines.
 *
 * @author R. S. Doiel, <rsdoiel@yahoo.com>
 * copyright (c) 2014
 * All rights reserved.
 * @license BSD 2-Clause License
 */
package main

import (
    "./fsengine"
	"./ottoengine"
	"./wslog"
    "./app"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"flag"
	"fmt"
	"log"
	"math/big"
	"net"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"
)

var REVISION = "v0.0.0-alpha"

// variables for keygen
var (
	cli_keygen   = flag.Bool("keygen", false, "Generate TLS ceriticates and keys")
	cli_ssl_host = flag.String("keygen-ssl-host", "", "Comma-separated hostnames and IPs to generate a certificate for")
	validFrom    = flag.String("keygen-start-date", "", "Creation date formatted as Jan 1 15:04:05 2011")
	validFor     = flag.Duration("keygen-duration", 365*24*time.Hour, "Duration that certificate is valid for")
	organization = flag.String("keygen-organization", "Acme Co.", "Organization used to sign certificate")
	isCA         = flag.Bool("keygen-ca", false, "whether this cert should be its own Certificate Authority")
	rsaBits      = flag.Int("keygen-rsa-bits", 2048, "Size of RSA key to generate")
)

// command line parameters that override environment variables
var (
	cli_use_tls   *bool
	cli_docroot   *string
	cli_host      *string
	cli_port      *string
	cli_cert      *string
	cli_key       *string
	cli_otto      *bool
	cli_otto_path *string
	cli_version   *bool
)

var Usage = func() {
	flag.PrintDefaults()
}

func request_log(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		wslog.LogRequest(r.Method, r.URL, r.RemoteAddr, r.Proto, r.Referer(), r.UserAgent())
		handler.ServeHTTP(w, r)
	})
}

func Webserver(profile *app.Profile) error {
	// If otto is enabled add routes and handle them.
	if profile.Otto == true {
		otto_path, err := filepath.Abs(profile.Otto_Path)
		if err != nil {
			log.Fatalf("Can't read %s: %s\n", profile.Otto_Path, err)
		}
		programs, err := ottoengine.Load(otto_path)
		if err != nil {
			log.Fatalf("Load error: %s\n", err)
		}
		ottoengine.AddRoutes(programs)
	}

	// Restricted FileService excluding dot files and directories
	http.HandleFunc("/", func (w http.ResponseWriter, r *http.Request) {
        // hande off this request/response pair to the fsengine
        fsengine.Engine(profile, w, r)
    })

	// Now start up the server and log transactions
	if profile.Use_TLS == true {
		if profile.Cert == "" || profile.Key == "" {
			log.Fatalf("TLS set true but missing key or certificate")
		}
		log.Println("Starting https://" + net.JoinHostPort(profile.Hostname, profile.Port))
		return http.ListenAndServeTLS(net.JoinHostPort(profile.Hostname, profile.Port), profile.Cert, profile.Key, request_log(http.DefaultServeMux))
	}
	log.Println("Starting http://" + net.JoinHostPort(profile.Hostname, profile.Port))
	// Now start up the server and log transactions
	return http.ListenAndServe(net.JoinHostPort(profile.Hostname, profile.Port), request_log(http.DefaultServeMux))
}

func keygen(profile *app.Profile) error {
	home := os.Getenv("HOME")
	certFilename := profile.Cert
	if certFilename == "" {
		certFilename = path.Join(home, "etc/ws/cert.pem")
	}
	keyFilename := profile.Key
	if keyFilename == "" {
		keyFilename = path.Join(home, "etc/ws/key.pem")
	}

	hostnames := profile.Hostname
	if *cli_ssl_host != "" {
		hostnames = *cli_ssl_host
	}
	if hostnames == "" {
		log.Fatalf("Missing required -ssl-host parameter")
	}

	log.Printf("\n\n"+
		" Cert: %s\n"+
		"  Key: %s\n"+
		" Host: %s\n"+
		" Organization: %v\n"+
		"\n\n",
		certFilename,
		keyFilename,
		hostnames,
		*organization)

	priv, err := rsa.GenerateKey(rand.Reader, *rsaBits)
	if err != nil {
		log.Fatalf("failed to generate private key: %s", err)
	}

	var notBefore time.Time
	if len(*validFrom) == 0 {
		notBefore = time.Now()
	} else {
		notBefore, err = time.Parse("Jan 2 15:04:05 2006", *validFrom)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to parse creation date: %s\n", err)
			return err
		}
	}

	notAfter := notBefore.Add(*validFor)

	// end of ASN.1 time
	endOfTime := time.Date(2049, 12, 31, 2, 59, 59, 0, time.UTC)
	if notAfter.After(endOfTime) {
		notAfter = endOfTime
	}

	template := x509.Certificate{
		SerialNumber: new(big.Int).SetInt64(0),
		Subject: pkix.Name{
			Organization: []string{*organization},
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	hosts := strings.Split(hostnames, ",")
	for _, h := range hosts {
		if ip := net.ParseIP(h); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, h)
		}
	}

	if *isCA {
		template.IsCA = true
		template.KeyUsage |= x509.KeyUsageCertSign
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		log.Fatalf("Failed to create certificate: %s", err)
	}

	certOut, err := os.Create(certFilename)
	if err != nil {
		log.Fatalf("failed to open cert.pem for writing: %s", err)
	}
	pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	certOut.Close()
	log.Printf("written %s\n", certFilename)

	keyOut, err := os.OpenFile(keyFilename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		log.Print("failed to open key.pem for writing:", err)
		return err
	}
	pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})
	keyOut.Close()
	log.Printf("written %s\n", keyFilename)
	// We got this for so no errors
	return nil
}

func init() {
	// command line parameters that override environment variables, many have short forms too.
	shortform := " (short form)"

	// No short form
	cli_use_tls = flag.Bool("tls", false, "Turn on TLS (https) support with true, off with false (default is false)")
	cli_key = flag.String("key", "", "path to your SSL key pem file.")
	cli_cert = flag.String("cert", "", "path to your SSL cert pem file.")

	// These have short forms too
	msg := "document root"
	cli_docroot = flag.String("docroot", "", msg)
	flag.StringVar(cli_docroot, "D", "", msg+shortform)

	msg = "hostname for webserver"
	cli_host = flag.String("host", "", msg)
	flag.StringVar(cli_host, "H", "", msg+shortform)

	msg = "Port number to listen on"
	cli_port = flag.String("port", "", msg)
	flag.StringVar(cli_port, "P", "", msg+shortform)

	msg = "turn on ottoengine, defaults to false"
	cli_otto = flag.Bool("otto", false, msg)
	flag.BoolVar(cli_otto, "o", false, msg+shortform)

	msg = "directory containingo your ottoengine JavaScript files"
	cli_otto_path = flag.String("otto-path", "", msg)
	flag.StringVar(cli_otto_path, "op", "", msg+shortform)

	msg = "Display the version number"
	cli_version = flag.Bool("version", false, msg)
	flag.BoolVar(cli_version, "v", false, msg+shortform)
}

func main() {
	flag.Parse()

	if *cli_version == true {
		fmt.Println(REVISION)
		os.Exit(0)
	}

	profile, _ := app.LoadProfile(*cli_docroot, *cli_host, *cli_port, *cli_use_tls, *cli_cert, *cli_key, *cli_otto, *cli_otto_path)
	if *cli_keygen == true {
		err := keygen(profile)
		if err != nil {
			log.Fatalf("%s\n", err)
		}
		os.Exit(0)
	}

	log.Printf("\n\n"+
		"          TLS: %t\n"+
		"         Cert: %s\n"+
		"          Key: %s\n"+
		"      Docroot: %s\n"+
		"         Host: %s\n"+
		"         Port: %s\n"+
		"       Run as: %s\n\n"+
		" Otto enabled: %t\n"+
		"         Path: %s\n"+
		"\n\n",
		profile.Use_TLS,
		profile.Cert,
		profile.Key,
		profile.Docroot,
		profile.Hostname,
		profile.Port,
		profile.Username,
		profile.Otto,
		profile.Otto_Path)
	err := Webserver(profile)
	if err != nil {
		log.Fatal(err)
	}
	os.Exit(0)
}
<|endoftext|>


<commit_before>package packer

import (
	"github.com/mitchellh/iochan"
	"io"
	"strings"
	"time"
)

// RemoteCmd represents a remote command being prepared or run.
type RemoteCmd struct {
	// Command is the command to run remotely. This is executed as if
	// it were a shell command, so you are expected to do any shell escaping
	// necessary.
	Command string

	// Stdin specifies the process's standard input. If Stdin is
	// nil, the process reads from an empty bytes.Buffer.
	Stdin io.Reader

	// Stdout and Stderr represent the process's standard output and
	// error.
	//
	// If either is nil, it will be set to ioutil.Discard.
	Stdout io.Writer
	Stderr io.Writer

	// This will be set to true when the remote command has exited. It
	// shouldn't be set manually by the user, but there is no harm in
	// doing so.
	Exited bool

	// Once Exited is true, this will contain the exit code of the process.
	ExitStatus int
}

// A Communicator is the interface used to communicate with the machine
// that exists that will eventually be packaged into an image. Communicators
// allow you to execute remote commands, upload files, etc.
//
// Communicators must be safe for concurrency, meaning multiple calls to
// Start or any other method may be called at the same time.
type Communicator interface {
	// Start takes a RemoteCmd and starts it. The RemoteCmd must not be
	// modified after being used with Start, and it must not be used with
	// Start again. The Start method returns immediately once the command
	// is started. It does not wait for the command to complete. The
	// RemoteCmd.Exited field should be used for this.
	Start(*RemoteCmd) error

	// Upload uploads a file to the machine to the given path with the
	// contents coming from the given reader. This method will block until
	// it completes.
	Upload(string, io.Reader) error

	// Download downloads a file from the machine from the given remote path
	// with the contents writing to the given writer. This method will
	// block until it completes.
	Download(string, io.Writer) error
}

// StartWithUi runs the remote command and streams the output to any
// configured Writers for stdout/stderr, while also writing each line
// as it comes to a Ui.
func (r *RemoteCmd) StartWithUi(c Communicator, ui Ui) error {
	stdout_r, stdout_w := io.Pipe()
	stderr_r, stderr_w := io.Pipe()

	// Set the writers for the output so that we get it streamed to us
	if r.Stdout == nil {
		r.Stdout = stdout_w
	} else {
		r.Stdout = io.MultiWriter(r.Stdout, stdout_w)
	}

	if r.Stderr == nil {
		r.Stderr = stderr_w
	} else {
		r.Stderr = io.MultiWriter(r.Stderr, stderr_w)
	}

	// Start the command
	if err := c.Start(r); err != nil {
		return err
	}

	// Create the channels we'll use for data
	exitCh := make(chan int, 1)
	stdoutCh := iochan.DelimReader(stdout_r, '\n')
	stderrCh := iochan.DelimReader(stderr_r, '\n')

	// Start the goroutine to watch for the exit
	go func() {
		defer stdout_w.Close()
		defer stderr_w.Close()
		r.Wait()
		exitCh <- r.ExitStatus
	}()

	// Loop and get all our output
OutputLoop:
	for {
		select {
		case output := <-stderrCh:
			ui.Message(strings.TrimSpace(output))
		case output := <-stdoutCh:
			ui.Message(strings.TrimSpace(output))
		case <-exitCh:
			break OutputLoop
		}
	}

	// Make sure we finish off stdout/stderr because we may have gotten
	// a message from the exit channel before finishing these first.
	for output := range stdoutCh {
		ui.Message(strings.TrimSpace(output))
	}

	for output := range stderrCh {
		ui.Message(strings.TrimSpace(output))
	}

	return nil
}

// Wait waits for the remote command to complete.
func (r *RemoteCmd) Wait() {
	for !r.Exited {
		time.Sleep(50 * time.Millisecond)
	}
}
<commit_msg>packer: In the case of an error, close stdout/stderr writers<commit_after>package packer

import (
	"github.com/mitchellh/iochan"
	"io"
	"strings"
	"time"
)

// RemoteCmd represents a remote command being prepared or run.
type RemoteCmd struct {
	// Command is the command to run remotely. This is executed as if
	// it were a shell command, so you are expected to do any shell escaping
	// necessary.
	Command string

	// Stdin specifies the process's standard input. If Stdin is
	// nil, the process reads from an empty bytes.Buffer.
	Stdin io.Reader

	// Stdout and Stderr represent the process's standard output and
	// error.
	//
	// If either is nil, it will be set to ioutil.Discard.
	Stdout io.Writer
	Stderr io.Writer

	// This will be set to true when the remote command has exited. It
	// shouldn't be set manually by the user, but there is no harm in
	// doing so.
	Exited bool

	// Once Exited is true, this will contain the exit code of the process.
	ExitStatus int
}

// A Communicator is the interface used to communicate with the machine
// that exists that will eventually be packaged into an image. Communicators
// allow you to execute remote commands, upload files, etc.
//
// Communicators must be safe for concurrency, meaning multiple calls to
// Start or any other method may be called at the same time.
type Communicator interface {
	// Start takes a RemoteCmd and starts it. The RemoteCmd must not be
	// modified after being used with Start, and it must not be used with
	// Start again. The Start method returns immediately once the command
	// is started. It does not wait for the command to complete. The
	// RemoteCmd.Exited field should be used for this.
	Start(*RemoteCmd) error

	// Upload uploads a file to the machine to the given path with the
	// contents coming from the given reader. This method will block until
	// it completes.
	Upload(string, io.Reader) error

	// Download downloads a file from the machine from the given remote path
	// with the contents writing to the given writer. This method will
	// block until it completes.
	Download(string, io.Writer) error
}

// StartWithUi runs the remote command and streams the output to any
// configured Writers for stdout/stderr, while also writing each line
// as it comes to a Ui.
func (r *RemoteCmd) StartWithUi(c Communicator, ui Ui) error {
	stdout_r, stdout_w := io.Pipe()
	stderr_r, stderr_w := io.Pipe()
	defer stdout_w.Close()
	defer stderr_w.Close()

	// Set the writers for the output so that we get it streamed to us
	if r.Stdout == nil {
		r.Stdout = stdout_w
	} else {
		r.Stdout = io.MultiWriter(r.Stdout, stdout_w)
	}

	if r.Stderr == nil {
		r.Stderr = stderr_w
	} else {
		r.Stderr = io.MultiWriter(r.Stderr, stderr_w)
	}

	// Start the command
	if err := c.Start(r); err != nil {
		return err
	}

	// Create the channels we'll use for data
	exitCh := make(chan int, 1)
	stdoutCh := iochan.DelimReader(stdout_r, '\n')
	stderrCh := iochan.DelimReader(stderr_r, '\n')

	// Start the goroutine to watch for the exit
	go func() {
		defer stdout_w.Close()
		defer stderr_w.Close()
		r.Wait()
		exitCh <- r.ExitStatus
	}()

	// Loop and get all our output
OutputLoop:
	for {
		select {
		case output := <-stderrCh:
			ui.Message(strings.TrimSpace(output))
		case output := <-stdoutCh:
			ui.Message(strings.TrimSpace(output))
		case <-exitCh:
			break OutputLoop
		}
	}

	// Make sure we finish off stdout/stderr because we may have gotten
	// a message from the exit channel before finishing these first.
	for output := range stdoutCh {
		ui.Message(strings.TrimSpace(output))
	}

	for output := range stderrCh {
		ui.Message(strings.TrimSpace(output))
	}

	return nil
}

// Wait waits for the remote command to complete.
func (r *RemoteCmd) Wait() {
	for !r.Exited {
		time.Sleep(50 * time.Millisecond)
	}
}
<|endoftext|>


<commit_before>package flags

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/lib"
)

// AccountNameFlag is used for all --account flags, excluding the global one.
type AccountNameFlag struct {
	// AccountName is the
	AccountName string
	// Value is the raw input to the flag, and can be used as the default when
	// creating the flag.
	Value string
	// SetFromCommandLine is false by default but is set to true when Set is
	// called. This allows setting a default by setting Value by yourself - Set
	// is called from urfave/cli's flag-parsing code.
	SetFromCommandLine bool
}

// Set sets Value and SetFromCommandLine on the flag
func (name *AccountNameFlag) Set(value string) error {
	name.Value = value
	name.SetFromCommandLine = true
	return nil
}

// Preprocess sets the value of this flag to the global account flag if it's
// unset, and then runs lib.ParseAccountName to set AccountName. This is an
// implementation of `app.Preprocessor`, which is detected and called
// automatically by actions created with `app.Action`
func (name *AccountNameFlag) Preprocess(c *app.Context) (err error) {
	if name.Value == "" {
		name.Value = c.Context.GlobalString("account")
	}
	name.AccountName = lib.ParseAccountName(name.Value, c.Config().GetIgnoreErr("account"))
	return
}

// String returns the AccountNameFlag as a string.
func (name AccountNameFlag) String() string {
	if name.AccountName == "" {
		return name.Value
	}
	return name.AccountName
}
<commit_msg>fix documentation comment for AccountName<commit_after>package flags

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/lib"
)

// AccountNameFlag is used for all --account flags, excluding the global one.
type AccountNameFlag struct {
	// AccountName is the actual name that will be passed on to API calls, and
	// is made by checking the contents of Value are a valid account. If Value
	// is unset then the value of the 'account' config variable is used
	AccountName string
	// Value is the raw input to the flag, and can be used as the default when
	// creating the flag.
	Value string
	// SetFromCommandLine is false by default but is set to true when Set is
	// called. This allows setting a default by setting Value by yourself - Set
	// is called from urfave/cli's flag-parsing code.
	SetFromCommandLine bool
}

// Set sets Value and SetFromCommandLine on the flag
func (name *AccountNameFlag) Set(value string) error {
	name.Value = value
	name.SetFromCommandLine = true
	return nil
}

// Preprocess sets the value of this flag to the global account flag if it's
// unset, and then runs lib.ParseAccountName to set AccountName. This is an
// implementation of `app.Preprocessor`, which is detected and called
// automatically by actions created with `app.Action`
func (name *AccountNameFlag) Preprocess(c *app.Context) (err error) {
	if name.Value == "" {
		name.Value = c.Context.GlobalString("account")
	}
	name.AccountName = lib.ParseAccountName(name.Value, c.Config().GetIgnoreErr("account"))
	return
}

// String returns the AccountNameFlag as a string.
func (name AccountNameFlag) String() string {
	if name.AccountName == "" {
		return name.Value
	}
	return name.AccountName
}
<|endoftext|>