github.com/diamondburned/arikawa/v2@v2.1.0/session/session.go (about)

     1  // Package session abstracts around the REST API and the Gateway, managing both
     2  // at once. It offers a handler interface similar to that in discordgo for
     3  // Gateway events.
     4  package session
     5  
     6  import (
     7  	"context"
     8  
     9  	"github.com/pkg/errors"
    10  
    11  	"github.com/diamondburned/arikawa/v2/api"
    12  	"github.com/diamondburned/arikawa/v2/gateway"
    13  	"github.com/diamondburned/arikawa/v2/internal/handleloop"
    14  	"github.com/diamondburned/arikawa/v2/utils/handler"
    15  )
    16  
    17  var ErrMFA = errors.New("account has 2FA enabled")
    18  
    19  // Closed is an event that's sent to Session's command handler. This works by
    20  // using (*Gateway).AfterClose. If the user sets this callback, no Closed events
    21  // would be sent.
    22  //
    23  // Usage
    24  //
    25  //    ses.AddHandler(func(*session.Closed) {})
    26  //
    27  type Closed struct {
    28  	Error error
    29  }
    30  
    31  // Session manages both the API and Gateway. As such, Session inherits all of
    32  // API's methods, as well has the Handler used for Gateway.
    33  type Session struct {
    34  	*api.Client
    35  	Gateway *gateway.Gateway
    36  
    37  	// Command handler with inherited methods.
    38  	*handler.Handler
    39  
    40  	// internal state to not be copied around.
    41  	looper *handleloop.Loop
    42  }
    43  
    44  func NewWithIntents(token string, intents ...gateway.Intents) (*Session, error) {
    45  	g, err := gateway.NewGatewayWithIntents(token, intents...)
    46  	if err != nil {
    47  		return nil, errors.Wrap(err, "failed to connect to Gateway")
    48  	}
    49  
    50  	return NewWithGateway(g), nil
    51  }
    52  
    53  // New creates a new session from a given token. Most bots should be using
    54  // NewWithIntents instead.
    55  func New(token string) (*Session, error) {
    56  	// Create a gateway
    57  	g, err := gateway.NewGateway(token)
    58  	if err != nil {
    59  		return nil, errors.Wrap(err, "failed to connect to Gateway")
    60  	}
    61  
    62  	return NewWithGateway(g), nil
    63  }
    64  
    65  // Login tries to log in as a normal user account; MFA is optional.
    66  func Login(email, password, mfa string) (*Session, error) {
    67  	// Make a scratch HTTP client without a token
    68  	client := api.NewClient("")
    69  
    70  	// Try to login without TOTP
    71  	l, err := client.Login(email, password)
    72  	if err != nil {
    73  		return nil, errors.Wrap(err, "failed to login")
    74  	}
    75  
    76  	if l.Token != "" && !l.MFA {
    77  		// We got the token, return with a new Session.
    78  		return New(l.Token)
    79  	}
    80  
    81  	// Discord requests MFA, so we need the MFA token.
    82  	if mfa == "" {
    83  		return nil, ErrMFA
    84  	}
    85  
    86  	// Retry logging in with a 2FA token
    87  	l, err = client.TOTP(mfa, l.Ticket)
    88  	if err != nil {
    89  		return nil, errors.Wrap(err, "failed to login with 2FA")
    90  	}
    91  
    92  	return New(l.Token)
    93  }
    94  
    95  func NewWithGateway(gw *gateway.Gateway) *Session {
    96  	handler := handler.New()
    97  	looper := handleloop.NewLoop(handler)
    98  
    99  	return &Session{
   100  		Gateway: gw,
   101  		// Nab off gateway's token
   102  		Client:  api.NewClient(gw.Identifier.Token),
   103  		Handler: handler,
   104  		looper:  looper,
   105  	}
   106  }
   107  
   108  func (s *Session) Open() error {
   109  	// Start the handler beforehand so no events are missed.
   110  	s.looper.Start(s.Gateway.Events)
   111  
   112  	// Set the AfterClose's handler.
   113  	s.Gateway.AfterClose = func(err error) {
   114  		s.Handler.Call(&Closed{
   115  			Error: err,
   116  		})
   117  	}
   118  
   119  	if err := s.Gateway.Open(); err != nil {
   120  		return errors.Wrap(err, "failed to start gateway")
   121  	}
   122  
   123  	return nil
   124  }
   125  
   126  // WithContext returns a shallow copy of Session with the context replaced in
   127  // the API client. All methods called on the returned Session will use this
   128  // given context.
   129  //
   130  // This method is thread-safe only after Open and before Close are called. Open
   131  // and Close should not be called on the returned Session.
   132  func (s *Session) WithContext(ctx context.Context) *Session {
   133  	cpy := *s
   134  	cpy.Client = s.Client.WithContext(ctx)
   135  	return &cpy
   136  }
   137  
   138  // Close closes the gateway. The connection is still resumable with the given
   139  // session ID.
   140  func (s *Session) Close() error {
   141  	return s.close(false)
   142  }
   143  
   144  // CloseGracefully permanently closes the gateway. The session ID is invalidated
   145  // afterwards.
   146  func (s *Session) CloseGracefully() error {
   147  	return s.close(true)
   148  }
   149  
   150  func (s *Session) close(gracefully bool) error {
   151  	// Stop the event handler
   152  	s.looper.Stop()
   153  	// Close the websocket
   154  	if gracefully {
   155  		return s.Gateway.CloseGracefully()
   156  	}
   157  	return s.Gateway.Close()
   158  }