package rawhttp.cookies.persist;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.net.CookieManager;
import java.net.CookieStore;
import java.net.HttpCookie;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
/**
* An implementation of {@link CookieStore} that persists cookies in a file.
* <p>
* It follows logic similar to a browser:
*
* <ul>
* <li>only persistent cookies are written out to the file (non-persistent cookies are only stored in memory).</li>
* <li>expired cookies are eventually deleted, but never returned when queried.</li>
* </ul>
* <p>
* How often cookies are flushed to the file depends on the implementation of {@link FlushStrategy} used by this
* cookie jar. By default, a {@link JvmShutdownFlushStrategy} is used.
*/
public class FileCookieJar implements CookieStore {
/**
* Strategy for how a {@link FileCookieJar} should flush cookies to the file.
*/
public interface FlushStrategy {
/**
* Initialize this strategy.
* <p>
* The given {@code flush} callable should be called every time the cookies should be flushed to a file.
*
* @param flush a callable that should be called to flush the cookies
*/
void init(Callable<Integer> flush);
/**
* This method is called every time the {@link FileCookieJar} is modified.
*
* @param cookieStore a view of the in-memory cookie store backing the {@link FileCookieJar}.
*/
void onUpdate(CookieStore cookieStore);
}
private final File file;
private final FlushStrategy flushPolicy;
private final Map<URI, List<HttpCookie>> persistentCookies = new HashMap<>();
// the InMemory implementation is not exposed, but we can steal it by creating a CookieManager like this:
private final CookieStore inMemory = new CookieManager().getCookieStore();
public FileCookieJar(File file) throws IOException {
this(file, new JvmShutdownFlushStrategy());
}
public FileCookieJar(File file, FlushStrategy flushPolicy) throws IOException {
this.file = file;
this.flushPolicy = flushPolicy;
flushPolicy.init(this::flush);
load();
}
public File getFile() {
return file;
}
public FlushStrategy getFlushPolicy() {
return flushPolicy;
}
private void load() throws IOException {
if (!file.exists()) {
return;
}
URI currentURI = null;
for (String line : Files.readAllLines(file.toPath(), StandardCharsets.UTF_8)) {
if (line.startsWith(" ")) {
if (currentURI == null) {
throw new IllegalStateException("Invalid syntax: cookie line found before a matching URI");
}
HttpCookie cookie = readCookie(line);
if (cookie != null) {
add(currentURI, cookie);
}
} else {
currentURI = URI.create(line);
}
}
}
private int flush() throws IOException {
int count = 0;
try (FileWriter writer = new FileWriter(file)) {
for (Map.Entry<URI, List<HttpCookie>> entry : persistentCookies.entrySet()) {
URI uri = entry.getKey();
writer.write(uri.toString());
writer.write('\n');
for (HttpCookie httpCookie : entry.getValue()) {
long maxAge = httpCookie.getMaxAge();
if (maxAge > 0 && !httpCookie.getDiscard()) {
long expiresAt = maxAge + System.currentTimeMillis() / 1000L;
writer.write(' ');
writeCookie(writer, httpCookie, expiresAt);
writer.write('\n');
count++;
}
}
}
}
return count;
}
private static HttpCookie readCookie(String line) {
String[] parts = line.split("\"");
assert parts.length == 20;
long maxAge = Long.parseLong(parts[19]) - System.currentTimeMillis() / 1000L;
if (maxAge < 0L) return null;
HttpCookie cookie = new HttpCookie(parts[1], parts[3]);
if (!parts[5].isEmpty()) cookie.setDomain(parts[5]);
if (!parts[7].isEmpty()) cookie.setPath(parts[7]);
if (!parts[9].isEmpty()) cookie.setComment(parts[9]);
if (!parts[11].isEmpty()) cookie.setCommentURL(parts[11]);
if (!parts[13].isEmpty()) cookie.setPortlist(parts[13]);
cookie.setVersion(Integer.parseInt(parts[15]));
cookie.setSecure(Boolean.parseBoolean(parts[17]));
cookie.setMaxAge(maxAge);
return cookie;
}
private static void writeCookie(FileWriter writer, HttpCookie httpCookie, long expiresAt) throws IOException {
writer.write(persistentValue(httpCookie.getName(), true));
writer.write(persistentValue(httpCookie.getValue(), true));
writer.write(persistentValue(httpCookie.getDomain(), true));
writer.write(persistentValue(httpCookie.getPath(), true));
writer.write(persistentValue(httpCookie.getComment(), true));
writer.write(persistentValue(httpCookie.getCommentURL(), true));
writer.write(persistentValue(httpCookie.getPortlist(), true));
writer.write(persistentValue(httpCookie.getVersion(), true));
writer.write(persistentValue(httpCookie.getSecure(), true));
writer.write(persistentValue(expiresAt, false));
}
private static String persistentValue(Object obj, boolean trailingSpace) {
String value;
if (obj == null) value = "\"\"";
else value = "\"" + obj + "\"";
if (trailingSpace) return value + ' ';
return value;
}
///////// Write methods /////////
@Override
public void add(URI uri, HttpCookie cookie) {
inMemory.add(uri, cookie);
if (cookie.getMaxAge() > 0) {
persistentCookies.computeIfAbsent(uri, (u) -> new ArrayList<>(4)).add(cookie);
flushPolicy.onUpdate(inMemory);
}
}
@Override
public boolean remove(URI uri, HttpCookie cookie) {
boolean result = inMemory.remove(uri, cookie);
List<HttpCookie> stored = persistentCookies.get(uri);
if (stored != null && stored.remove(cookie)) {
flushPolicy.onUpdate(inMemory);
if (stored.isEmpty()) {
persistentCookies.remove(uri);
}
}
return result;
}
@Override
public boolean removeAll() {
boolean result = inMemory.removeAll();
if (!persistentCookies.isEmpty()) {
persistentCookies.clear();
flushPolicy.onUpdate(inMemory);
}
return result;
}
///////// Read-only methods /////////
@Override
public List<HttpCookie> get(URI uri) {
return inMemory.get(uri);
}
@Override
public List<HttpCookie> getCookies() {
return inMemory.getCookies();
}
@Override
public List<URI> getURIs() {
return inMemory.getURIs();
}
}
