1   /*
2    * Copyright 2018 LINE Corporation
3    *
4    * LINE Corporation licenses this file to you under the Apache License,
5    * version 2.0 (the "License"); you may not use this file except in compliance
6    * with the License. You may obtain a copy of the License at:
7    *
8    *   https://www.apache.org/licenses/LICENSE-2.0
9    *
10   * Unless required by applicable law or agreed to in writing, software
11   * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12   * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13   * License for the specific language governing permissions and limitations
14   * under the License.
15   */
16  package com.linecorp.centraldogma.server.auth.shiro;
17  
18  import static com.google.common.base.Preconditions.checkArgument;
19  import static java.util.Objects.requireNonNull;
20  
21  import java.io.Serializable;
22  import java.time.Duration;
23  import java.util.Collection;
24  import java.util.function.Supplier;
25  
26  import org.apache.shiro.session.Session;
27  import org.apache.shiro.session.UnknownSessionException;
28  import org.apache.shiro.session.mgt.SimpleSession;
29  import org.apache.shiro.session.mgt.eis.MemorySessionDAO;
30  import org.apache.shiro.session.mgt.eis.SessionDAO;
31  
32  import com.github.benmanes.caffeine.cache.Cache;
33  import com.github.benmanes.caffeine.cache.Caffeine;
34  
35  /**
36   * A {@link SessionDAO} which stores {@link Session}s in the memory. Unlike {@link MemorySessionDAO},
37   * its memory usage is limited to the configured size.
38   */
39  final class LimitedMemorySessionDAO implements SessionDAO {
40  
41      private final Supplier<String> sessionIdGenerator;
42      private final Cache<Serializable, Session> cache;
43  
44      /**
45       * Creates a new instance.
46       *
47       * @param sessionIdGenerator the session ID generator
48       * @param maximumSize the maximum number of sessions can be stored in the cache
49       * @param maximumDuration the maximum duration which the session stays in the cache since entering
50       */
51      LimitedMemorySessionDAO(Supplier<String> sessionIdGenerator, int maximumSize, Duration maximumDuration) {
52          this.sessionIdGenerator = requireNonNull(sessionIdGenerator, "sessionIdGenerator");
53          checkArgument(maximumSize > 0, "maximumSize: %s (expected: > 0)", maximumSize);
54          requireNonNull(maximumDuration, "maximumDuration");
55          cache = Caffeine.newBuilder()
56                          .expireAfterWrite(maximumDuration)
57                          .maximumSize(maximumSize)
58                          .build();
59      }
60  
61      @Override
62      public Serializable create(Session session) {
63          final SimpleSession simpleSession = ensureSimpleSession(session);
64          final String id = sessionIdGenerator.get();
65          simpleSession.setId(id);
66          cache.put(id, simpleSession);
67          return session.getId();
68      }
69  
70      @Override
71      public Session readSession(Serializable sessionId) {
72          if (sessionId == null) {
73              throw new UnknownSessionException("sessionId is null");
74          }
75  
76          final Session session = cache.getIfPresent(sessionId);
77          if (session != null) {
78              return session;
79          }
80          throw new UnknownSessionException(sessionId.toString());
81      }
82  
83      @Override
84      public void update(Session session) {
85          final SimpleSession simpleSession = ensureSimpleSession(session);
86          readSession(simpleSession.getId());
87          cache.put(simpleSession.getId(), simpleSession);
88      }
89  
90      @Override
91      public void delete(Session session) {
92          cache.invalidate(session.getId());
93      }
94  
95      @Override
96      public Collection<Session> getActiveSessions() {
97          return cache.asMap().values();
98      }
99  
100     private static SimpleSession ensureSimpleSession(Session session) {
101         requireNonNull(session, "session");
102         checkArgument(session instanceof SimpleSession,
103                       "session: %s (expected: SimpleSession)", session);
104         return (SimpleSession) session;
105     }
106 }