TL;DR: and obligatory advice:
Don't do this. There are too many problems with this whole design. Even if the spec is much clearer in your head than it is in this description, you just don't understand what is involved in designing a system like this. Rolling your own cryptosystem isn't quite so bad an idea as rolling your own crypto primitives, but it's still a very bad idea unless you have an excellent understanding of applied cryptography. Cryptography is hard! For example, the SSL/TLS protocol that is used to secure everything from this website to international banking is currently on its sixth released version (and everything before the fifth is deprecated as unsafe), because even in widely-used protocols designed by experts, people keep finding problems with the old ones.
If you want to create yet another end-to-end encrypted (E2EE) messaging app, do yourself a favor and just adopt the excellent, and open-source, Signal protocol, or something descended from it and similarly well reviewed/audited. This protocol is already the basis of the Signal and WhatsApp messaging apps, optionally used by other apps like Skype and Facebook Messenger, and its design principles are the bases of some other apps (like Telegram and Wire) and protocols (like Matrix).
There are a bunch of problems, or at least confusing errors, with this scheme.
Send Public Key to Server (AWS)
Is there any way for a user to verify another user's public key? Or does that not even come up, in this system? If your encryption relies on the server in the middle being trustworthy, that can technically be end-to-end encryption but it misses the "zero-trust" property that is usually desired of E2EE systems.
Using another Key pair server will encrypt and store public key
Why encrypt the public key? It's... public. Literally everybody is allowed to know what it is, that's the point of public keys. Also when you say "another Key pair" do you mean one that the server has generated?
When Another user wants to send message server will provide encrypted public key and another user will encrypt with his private key
Two problems here.
First of all, you say "server will provide encrypted public key". Whose public key is that (the sender's? the server's? the recipient's?)? Whose key was used to encrypt it (and therefore can decrypt it)? And again, why is it encrypted at all?
Second, you say "another user will encrypt with his private key". I assume "another user" here means the person who wants to send a message? You need to be explicit about that kind of thing, there are at least two parties to every message and that's not counting the server. The real problem, though, is that "encrypt with his private key" is not a valid operation. You can't use private keys for encryption. You can only sign or decrypt with them. Did you mean that the sender will sign the message with their private key? Again, be very specific when designing or describing a system like this: which party (person/software) is doing what operation (encrypt/decrypt/sign/verify/transmit/store/etc.) on what object (message/some party's public key/some party's private key/etc.) with what other relevant details (destination/metadata/transmission channel/etc.).
Server will decrypt that message with own key pair
Why is the server decrypting the message? Why is it even encrypted to the server's public key? That completely violates end-to-end encryption. You can't have E2EE unless the message not only is not but can not be decrypted by any party except the final recipient.
Server sends updated message to first message and he will decrypt with his private key
Two problems again.
"Server sends updated message to first message" what does this mean? You can't send a message to a message, messages aren't the kind of entity that can receive messages. Do you mean sending an updated version of the first message? (That just raises more questions, like what is the update and why is it needed and how is it encrypted and so on.) Do you mean sending a new (second) message to the recipient of the first message? (Most of the same questions.) Do you mean sending a message back from the server to the sender of the initial message? (Perhaps to tell them the message was delivered or something?) Again, you need to be clear and specific with these things.
"he will decrypt with his private key" Who is "he" here and what is being decrypted? Don't ever use pronouns in a spec like this, be specific (you can assign labels like "person X" and "person Y", or fake names like "Alice" and "Bob", but you need some specific way to refer to every entity - including every message and every key - when designing or describing a system, and you need to consistently use it).
Other questions about this whole scheme:
- Are users able to change private keys / rotate to new private keys? What does that take? Do other users know when that happens (necessary, if you're going to support out-of-band key verification)?
- What encryption are you using? You talk about encrypting with public keys, which implies you're using RSA for asymmetric key pairs, but RSA is incredibly slow and totally unsuited to encryption of arbitrary-length data, so it's approximately always used in a hybrid cryptosystem where the asymmetric key pair is used to securely exchange a symmetric key (usually AES or a stream cipher like ChaCha20) and the symmetric cipher (with appropriate nonces/IVs, plus mode of operation and padding if needed) is used for bulk encryption and decryption of messages.
- Are you providing message integrity? Encryption doesn't prevent an attacker from modifying the messages in transit, and if the attacker already knows at least part of the message, it is often possible to precisely change that part to say something else. Options for integrity include message authentication codes (e.g. HMAC), authenticated encryption modes (e.g. GCM as in AES-GCM, or Poly1305 with ChaCha20), or asymmetric digital signatures.
- Do you care about non-repudiation? If it's important to you that each message be cryptographically verifiable to come from a specific source (such that there's no way, after the fact, to deny sending a.k.a. repudiate a message you did in fact send) then you'll want to include digital signatures on each message (and expose them to the recipients, so that the sender can't repudiate the message). If you instead want to support repudiation, you'll need to not use digital signatures (HMACs or authenticated encryption are fine, since those use symmetric keys and it's impossible, after the fact, to know which party to the communication created or modified a given message).
- Do you want to provide forward secrecy (the property that even if a private key is later compromised, the attacker can't use it to decrypt previously-recorded conversations)? If so, you'll need to either only support online real-time communication, or you'll need to use something like the double-ratchet algorithm that Signal protocol is based on.
- Do you want to support group messaging? If so, there's a whole new raft of complicated questions, like handling people joining and leaving groups, or the performance of communicating with a large group. Or you can use a protocol that has already solved that.
- Do you want the app to store message history in plain text, encrypted with a client-side key, in the original encrypted form, or not at all? Each has its own implications and complications. Note that, aside from if you want "original encrypted form", this has no implications for protocol design but is an important app design consideration. (Original encrypted form is generally a bad idea, for reasons including "can't use ephemeral keys"; you can get forward secrecy for the user's private key but the session key needs to persist.)
- What are you using for message transport? Are you sending encrypted messages over a plain-text transport or using something like TLS?
There's more, but it really depends on the context of this app (who its audience is, where and how it's going to be used, what its threat model is, etc.) to even know what questions to ask.
Your spec doesn't say anything about "logging out" and "logging in", but I assume you mean that there's some kind of credentials (perhaps username+password or similar) that each user uses? Typically, messaging systems like this have a few ways to handle private keys, depending on how much they trust the server (ranging from "fully, don't need E2EE" through "want E2EE but don't need zero-trust" to "not at all, zero trust").
- Store both public and private keys on the server. This means the server can decrypt anybody's message, whether it does so in person or not, so you can't guarantee E2EE and definitely don't have zero-trust.
- Store both public and encrypted private keys on the server, and encrypt each user's private key with a key derived (from the user's password, usually) on the client, such that the server never sees the plain-text private key and can't decrypt the private keys it holds. This is compatible with zero-trust but requires a complicated auth scheme, since if you're deriving keys from the user's password then that password can't itself ever be exposed to the server. If you don't need zero-trust, this works using conventional password authentication but a malicious or compromised server would be able to decrypt the private keys and break the E2EE.
- Store only public keys on the server. Private keys are moved from client to client directly, by the user, without ever touching the server (alternatively, encrypted private keys are passed through the server but the key-encrypting key is only passed directly from one client to another, rather than deriving it from user input as above). Compatible with (and a common method for) zero-trust E2EE, but makes it inconvenient to use a new client for the first time (user needs to have an existing client that they can transfer the key from, and there needs to be some method - often a QR code or a file - for doing the transfer).
Storing private keys on devices, or even web browsers, is generally fine. However, you should use each platform's key management features (like Keychain or SubtleCrypto) to store and use the private key (ideally in a non-retrievable way, if that's viable) rather than handling it directly. If you must handle it directly, try to store it in an encrypted format to prevent an attacker who gains access to the device storage from stealing the private key (though this might require the user to authenticate to the client via e.g. a password, each time the users wants to use the app).
